From 2b2c2030df560e07191f36fe5d1b820a64d2ec14 Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Thu, 27 Feb 2025 14:52:48 -0500
Subject: [PATCH 1/7] Allow running the model with SSGW (fully nonlinear wave)

---
 .github/workflows/build.yml |   2 +-
 test_twowave.py             |  48 +++-
 twowave.py                  | 421 ++++++++++++++++++++++++++++++++----
 3 files changed, 427 insertions(+), 44 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index eb3284e..156086d 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -20,7 +20,7 @@ jobs:
 
     - name: Install dependencies
       run: |
-        pip install -U pip pytest
+        pip install -U pip pytest git+https://github.com/wavesgroup/ssgw
         pip install -U .
 
     - name: Run tests
diff --git a/test_twowave.py b/test_twowave.py
index 14eeb8e..5a0640b 100644
--- a/test_twowave.py
+++ b/test_twowave.py
@@ -1,6 +1,12 @@
 import numpy as np
-from twowave import angular_frequency, elevation, gravity, WaveModulationModel
+from twowave import (
+    angular_frequency,
+    elevation,
+    gravity,
+    WaveModulationModel,
+)
 import xarray as xr
+import pytest
 
 
 G0 = 9.8
@@ -14,6 +20,27 @@ def test_angular_frequency():
 def test_elevation():
     assert elevation(0, 0, 1, 1, 1, wave_type="linear") == 1
 
+    # Create mock nonlinear properties
+    x = np.linspace(0, 2 * np.pi, 128)
+    z = 0.1 * np.cos(x)  # Simple cosine elevation
+    mock_props = (x, z, None, None, None, None)
+
+    # Test nonlinear elevation
+    assert np.isclose(
+        elevation(0, 0, 0.1, 1, 1, wave_type="nonlinear", nonlinear_props=mock_props),
+        0.1,
+        rtol=1e-3,
+    )
+
+    # Test nonlinear elevation at phase π
+    assert np.isclose(
+        elevation(
+            np.pi, 0, 0.1, 1, 1, wave_type="nonlinear", nonlinear_props=mock_props
+        ),
+        -0.1,
+        rtol=1e-3,
+    )
+
 
 def test_gravity():
     assert gravity(0, 0, 0, 1, G0, wave_type="linear") == G0
@@ -21,9 +48,28 @@ def test_gravity():
 
 
 def test_wave_modulation_model():
+    # Test with linear wave type
     m = WaveModulationModel(num_periods=1)
     m.run()
     ds = m.to_xarray()
     assert type(ds) == xr.Dataset
     assert np.all(np.isfinite(ds.wavenumber))
     assert np.all(np.isfinite(ds.amplitude))
+
+    # Test with nonlinear wave type
+    try:
+        m = WaveModulationModel(num_periods=1)
+        m.run(wave_type="nonlinear")
+        ds = m.to_xarray()
+        assert type(ds) == xr.Dataset
+        assert np.all(np.isfinite(ds.wavenumber))
+        assert np.all(np.isfinite(ds.amplitude))
+    except Exception as e:
+        # Skip this test if SSGW module is not available or other issues arise
+        pytest.skip(f"Skipping nonlinear test due to: {str(e)}")
+
+
+def test_missing_nonlinear_props():
+    # Test that appropriate error is raised when nonlinear_props is missing
+    with pytest.raises(ValueError, match="nonlinear_props must be provided"):
+        elevation(0, 0, 0.1, 1, 1, wave_type="nonlinear")
diff --git a/twowave.py b/twowave.py
index 75e502a..0b4fd35 100644
--- a/twowave.py
+++ b/twowave.py
@@ -4,6 +4,8 @@
 
 import numpy as np
 from rich.progress import track
+from scipy.interpolate import interp1d
+from ssgw import SSGW
 import xarray as xr
 
 
@@ -12,7 +14,13 @@ def angular_frequency(g: float, k: float, u: float = 0) -> float:
 
 
 def elevation(
-    x: float, t: float, a: float, k: float, omega: float, wave_type: str = "linear"
+    x: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
     phase = k * x - omega * t
     if wave_type == "linear":
@@ -22,25 +30,74 @@ def elevation(
         term2 = 0.5 * a * k * np.cos(2 * phase)
         term3 = (a * k) ** 2 * (3 / 8 * np.cos(3 * phase) - 1 / 16 * np.cos(phase))
         return a * (term1 + term2 + term3)
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, elevations, _, _, _, _ = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Interpolate elevation at the phase-shifted position
+            result[i] = np.interp(x_phase * k, positions, elevations, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
     else:
-        raise ValueError("wave_type must be either 'linear' or 'stokes'")
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def surface_slope(
-    x: float, t: float, a: float, k: float, omega: float, wave_type: str = "linear"
+    x: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
     phase = k * x - omega * t
     ak = a * k
     if wave_type == "linear":
         slope = -ak * np.sin(phase)
+        return slope
     elif wave_type == "stokes":
         term1 = -ak * np.sin(phase)
         term2 = -(ak**2) * np.sin(2 * phase)
         term3 = -(ak**3) * (9 / 8 * np.sin(3 * phase) - 1 / 16 * np.sin(phase))
         slope = term1 + term2 + term3
+        return slope
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, elevations, _, _, _, _ = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Calculate dense slopes once
+        dense_pos = np.linspace(0, 2 * np.pi, 1000)
+        dense_elev = np.interp(dense_pos, positions, elevations, period=2 * np.pi)
+        dx = dense_pos[1] - dense_pos[0]
+        slopes = np.gradient(dense_elev, dx)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Interpolate slope at the requested position
+            result[i] = np.interp(x_phase * k, dense_pos, slopes, period=2 * np.pi) / k
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
     else:
-        raise ValueError("wave_type must be either 'linear' or 'stokes'")
-    return slope
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def gravity(
@@ -51,18 +108,20 @@ def gravity(
     omega: float,
     g0: float = 9.8,
     wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
     """Gravitational acceleration at the surface of a long wave.
-    Supports both linear and Stokes wave types.
+    Supports linear, Stokes, and nonlinear wave types.
     """
     phase = k * x - omega * t
-    eta = elevation(x, t, a, k, omega, wave_type=wave_type)
 
     if wave_type == "linear":
+        eta = elevation(x, t, a, k, omega, wave_type)
         return g0 * (
             1 - a * k * np.exp(k * eta) * (np.cos(phase) - a * k * np.sin(phase) ** 2)
         )
     elif wave_type == "stokes":
+        eta = elevation(x, t, a, k, omega, wave_type)
         return g0 * (
             1
             - a
@@ -80,26 +139,110 @@ def gravity(
             )
             * np.exp(k * eta)
         )
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, _, _, _, _, az = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Get effective gravity at this position
+            result[i] = g0 + np.interp(x_phase * k, positions, az, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
     else:
-        raise ValueError("wave_type must be either 'linear' or 'stokes'")
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def orbital_horizontal_velocity(
-    x: float, z: float, t: float, a: float, k: float, omega: float
+    x: float,
+    z: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
-    """Horizontal orbital velocity at depth z."""
-    return a * omega * np.cos(k * x - omega * t) * np.exp(k * z)
+    """Horizontal orbital velocity at depth z.
+    Supports linear, Stokes, and nonlinear wave types.
+    """
+    if wave_type in ["linear", "stokes"]:
+        return a * omega * np.cos(k * x - omega * t) * np.exp(k * z)
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, _, u, _, _, _ = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Get horizontal velocity at this position
+            result[i] = np.interp(x_phase * k, positions, u, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
+    else:
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def orbital_vertical_velocity(
-    x: float, z: float, t: float, a: float, k: float, omega: float
+    x: float,
+    z: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
-    """Horizontal orbital velocity at depth z."""
-    return a * omega * np.sin(k * x - omega * t) * np.exp(k * z)
+    """Vertical orbital velocity at depth z.
+    Supports linear, Stokes, and nonlinear wave types.
+    """
+    if wave_type in ["linear", "stokes"]:
+        return a * omega * np.sin(k * x - omega * t) * np.exp(k * z)
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, _, _, w, _, _ = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Get vertical velocity at this position
+            result[i] = np.interp(x_phase * k, positions, w, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
+    else:
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def orbital_horizontal_acceleration(
-    x: float, t: float, a: float, k: float, omega: float, wave_type: str = "linear"
+    x: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
     phase = k * x - omega * t
     if wave_type == "linear":
@@ -107,6 +250,7 @@ def orbital_horizontal_acceleration(
         dU_dt = (
             a * omega**2 * np.exp(k * eta) * np.sin(phase) * (a * k * np.cos(phase) + 1)
         )
+        return dU_dt
     elif wave_type == "stokes":
         eta = elevation(x, t, a, k, omega, wave_type)
         term1 = a * omega**2 * np.exp(k * eta) * np.sin(phase)
@@ -124,13 +268,37 @@ def orbital_horizontal_acceleration(
             )
         )
         dU_dt = term1 + term2
+        return dU_dt
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, _, _, _, ax, _ = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Get horizontal acceleration at this position
+            result[i] = np.interp(x_phase * k, positions, ax, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
     else:
-        raise ValueError("wave_type must be either 'linear' or 'stokes'")
-    return dU_dt
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def orbital_vertical_acceleration(
-    x: float, t: float, a: float, k: float, omega: float, wave_type: str = "linear"
+    x: float,
+    t: float,
+    a: float,
+    k: float,
+    omega: float,
+    wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
     phase = k * x - omega * t
     if wave_type == "linear":
@@ -141,6 +309,7 @@ def orbital_vertical_acceleration(
             * np.exp(k * eta)
             * (a * k * np.sin(phase) ** 2 - np.cos(phase))
         )
+        return dW_dt
     elif wave_type == "stokes":
         eta = elevation(x, t, a, k, omega, wave_type)
         term1 = -a * omega**2 * np.exp(k * eta) * np.cos(phase)
@@ -158,9 +327,27 @@ def orbital_vertical_acceleration(
             )
         )
         dW_dt = term1 + term2
+        return dW_dt
+    elif wave_type == "nonlinear":
+        if nonlinear_props is None:
+            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
+        positions, _, _, _, _, az = nonlinear_props
+
+        # Handle both scalar and array inputs
+        x_array = np.atleast_1d(x)
+        result = np.zeros_like(x_array, dtype=float)
+
+        # Apply for each x value
+        for i, xi in enumerate(x_array):
+            # Apply phase shift based on time
+            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
+            # Get vertical acceleration at this position
+            result[i] = np.interp(x_phase * k, positions, az, period=2 * np.pi)
+
+        # Return scalar if input was scalar
+        return result[0] if np.isscalar(x) else result
     else:
-        raise ValueError("wave_type must be either 'linear' or 'stokes'")
-    return dW_dt
+        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
 def gravity_curvilinear(
@@ -171,14 +358,56 @@ def gravity_curvilinear(
     omega: float,
     g0: float = 9.8,
     wave_type: str = "linear",
+    nonlinear_props: tuple = None,
 ) -> float:
-    dU_dt = orbital_horizontal_acceleration(x, t, a, k, omega, wave_type)
-    dW_dt = orbital_vertical_acceleration(x, t, a, k, omega, wave_type)
-    slope = surface_slope(x, t, a, k, omega, wave_type)
+    dU_dt = orbital_horizontal_acceleration(
+        x, t, a, k, omega, wave_type, nonlinear_props
+    )
+    dW_dt = orbital_vertical_acceleration(x, t, a, k, omega, wave_type, nonlinear_props)
+    slope = surface_slope(x, t, a, k, omega, wave_type, nonlinear_props)
     g = g0 * np.cos(slope) + dW_dt * np.cos(slope) + dU_dt * np.sin(slope)
     return g
 
 
+def nonlinear_wave_properties(
+    a: float, k: float, g0: float = 9.8, num_points: int = 128
+):
+    ak = a * k
+    omega = np.sqrt(g0 * k)  # deep water
+    wave = SSGW(np.inf, ak, 128)
+    Cp = wave.ce * omega / k
+    x, z = wave.zs.real, wave.zs.imag
+    u, w = (wave.ws.real + wave.ce) * np.sqrt(g0), -wave.ws.imag * np.sqrt(g0)
+
+    # Shift the wave solution by pi
+    x -= np.pi
+    z = np.array(z[len(z) // 2 :].tolist() + z[: len(z) // 2].tolist())
+    u = np.array(u[len(u) // 2 :].tolist() + u[: len(u) // 2].tolist())
+    w = np.array(w[len(w) // 2 :].tolist() + w[: len(w) // 2].tolist())
+
+    dx = diff(x)
+    dx[dx < 0] += np.pi
+    a_z = -diff(w) / dx * Cp
+    a_x = -diff(u) / dx * Cp
+
+    # Shift back by pi
+    x += np.pi
+    z = np.array(z[len(z) // 2 :].tolist() + z[: len(z) // 2].tolist())
+    u = np.array(u[len(u) // 2 :].tolist() + u[: len(u) // 2].tolist())
+    w = np.array(w[len(w) // 2 :].tolist() + w[: len(w) // 2].tolist())
+    a_x = np.array(a_x[len(a_x) // 2 :].tolist() + a_x[: len(a_x) // 2].tolist())
+    a_z = np.array(a_z[len(a_z) // 2 :].tolist() + a_z[: len(a_z) // 2].tolist())
+
+    x_ = np.linspace(x[0], x[-1], num_points)
+    z = interp1d(x, z, kind="cubic")(x_)
+    u = interp1d(x, u, kind="cubic")(x_)
+    w = interp1d(x, w, kind="cubic")(x_)
+    a_x = interp1d(x, a_x, kind="cubic")(x_)
+    a_z = interp1d(x, a_z, kind="cubic")(x_)
+
+    return x_, z, u, w, a_x, a_z
+
+
 def diff(x: np.ndarray) -> np.ndarray:
     """2nd order, centered difference"""
     dx = np.zeros_like(x)
@@ -234,7 +463,7 @@ def __init__(
         grav0: float = 9.8,
         a_short: float | np.ndarray = 0.01,
         k_short: float | np.ndarray = 10,
-        grid_size: int = 100,
+        grid_size: int = 128,
         num_periods: int = 10,
         curvilinear: bool = True,
     ) -> None:
@@ -258,6 +487,9 @@ def __init__(
         self.num_time_steps = len(self.time)
         self.curvilinear = curvilinear
 
+        # Initialize nonlinear wave properties as None
+        self.nonlinear_props = None
+
     def get_elevation_ramp(self, t: float) -> float:
         """Determine the long-wave profile."""
         if self.ramp_type == None:
@@ -269,7 +501,7 @@ def get_elevation_ramp(self, t: float) -> float:
         elif self.ramp_type == "groups":
             group_duration = self.num_waves_in_group * self.T_long
             eta_ramp = np.sin(t / group_duration * np.pi) ** 2
-        return eta_ramp
+        return max(eta_ramp, 1e-6)
 
     def run(
         self,
@@ -280,13 +512,21 @@ def run(
         save_tendencies: bool = False,
     ):
         """Integrate the model forward in time."""
-        if wave_type not in ["linear", "stokes"]:
-            raise ValueError("Invalid wave_type")
+        if wave_type not in ["linear", "stokes", "nonlinear"]:
+            raise ValueError(
+                "Invalid wave_type, must be 'linear', 'stokes', or 'nonlinear'"
+            )
 
         self.elevation = elevation
-        self.gravity = gravity_curvilinear
+        self.gravity = gravity_curvilinear if self.curvilinear else gravity
         self._wave_type = wave_type
 
+        if wave_type == "nonlinear":
+            positions, elevations, u, w, ax, az = nonlinear_wave_properties(
+                self.a_long, self.k_long, self.grav0
+            )
+            self.nonlinear_props = (positions, elevations, u, w, ax, az)
+
         if not ramp_type in [None, "linear", "groups"]:
             raise ValueError("Invalid ramp_type")
 
@@ -307,6 +547,12 @@ def run(
         self.g = np.zeros_like(self.k)
         self.omega = np.zeros_like(self.k)
 
+        # Allocate and initialize surface velocities, elevation, and slope.
+        self.u = np.zeros_like(self.k)
+        self.w = np.zeros_like(self.k)
+        self.eta = np.zeros_like(self.k)
+        self.slope = np.zeros_like(self.k)
+
         # Allocate tendencies if requested.
         if self.save_tendencies:
             self.k_propagation_tendency = np.zeros_like(self.k)
@@ -347,28 +593,74 @@ def wavenumber_tendency(self, k, t):
         """Compute the tendencies of the wavenumber conservation balance at time t."""
         eta_ramp = self.get_elevation_ramp(t)
 
-        eta = eta_ramp * self.elevation(
-            self.x, t, self.a_long, self.k_long, self.omega_long, self._wave_type
+        if self._wave_type == "nonlinear":
+            self.nonlinear_props = nonlinear_wave_properties(
+                eta_ramp * self.a_long, self.k_long, self.grav0
+            )
+
+        eta = self.elevation(
+            self.x,
+            t,
+            eta_ramp * self.a_long,
+            self.k_long,
+            self.omega_long,
+            self._wave_type,
+            self.nonlinear_props,
         )
+        self.eta[self.current_time_step] = eta
 
         if self.curvilinear:
-            slope = eta_ramp * surface_slope(
-                self.x, t, self.a_long, self.k_long, self.omega_long, self._wave_type
+            slope = surface_slope(
+                self.x,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                self._wave_type,
+                self.nonlinear_props,
             )
+            self.slope[self.current_time_step] = slope
+
             alpha = np.arctan(slope)
             self.ds = self.dx / np.cos(alpha)
             u = orbital_horizontal_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
+            self.u[self.current_time_step] = u
+
             w = orbital_vertical_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
+            self.w[self.current_time_step] = w
+
             vel = u * np.cos(alpha) + w * np.sin(alpha)
         else:
             self.ds = self.dx * np.ones(self.grid_size)
             vel = orbital_horizontal_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
+            self.u[self.current_time_step] = vel
 
         g = self.gravity(
             self.x,
@@ -378,6 +670,7 @@ def wavenumber_tendency(self, k, t):
             self.omega_long,
             self.grav0,
             wave_type=self._wave_type,
+            nonlinear_props=self.nonlinear_props,
         )
         self.g[self.current_time_step] = g
 
@@ -408,27 +701,65 @@ def waveaction_tendency(self, N, t):
         """Compute the tendencies of the wave action balance at time t."""
         eta_ramp = self.get_elevation_ramp(t)
 
-        eta = eta_ramp * self.elevation(
-            self.x, t, self.a_long, self.k_long, self.omega_long, self._wave_type
+        if self._wave_type == "nonlinear":
+            self.nonlinear_props = nonlinear_wave_properties(
+                eta_ramp * self.a_long, self.k_long, self.grav0
+            )
+
+        eta = self.elevation(
+            self.x,
+            t,
+            eta_ramp * self.a_long,
+            self.k_long,
+            self.omega_long,
+            self._wave_type,
+            self.nonlinear_props,
         )
 
         if self.curvilinear:
-            slope = eta_ramp * surface_slope(
-                self.x, t, self.a_long, self.k_long, self.omega_long, self._wave_type
+            slope = surface_slope(
+                self.x,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                self._wave_type,
+                self.nonlinear_props,
             )
             alpha = np.arctan(slope)
             self.ds = self.dx / np.cos(alpha)
             u = orbital_horizontal_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
             w = orbital_vertical_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
             vel = u * np.cos(alpha) + w * np.sin(alpha)
         else:
             self.ds = self.dx * np.ones(self.grid_size)
             vel = orbital_horizontal_velocity(
-                self.x, eta, t, eta_ramp * self.a_long, self.k_long, self.omega_long
+                self.x,
+                eta,
+                t,
+                eta_ramp * self.a_long,
+                self.k_long,
+                self.omega_long,
+                wave_type=self._wave_type,
+                nonlinear_props=self.nonlinear_props,
             )
 
         g = self.gravity(
@@ -439,6 +770,7 @@ def waveaction_tendency(self, N, t):
             self.omega_long,
             self.grav0,
             wave_type=self._wave_type,
+            nonlinear_props=self.nonlinear_props,
         )
         Cg = (
             angular_frequency(g, self.k[self.current_time_step])
@@ -473,6 +805,11 @@ def to_xarray(self) -> xr.Dataset:
                 "wave_action": (("time", "space"), self.N),
                 "gravitational_acceleration": (("time", "space"), self.g),
                 "angular_frequency": (("time", "space"), self.omega),
+                "elevation": (("time", "space"), self.eta),
+                "slope": (("time", "space"), self.slope),
+                "horizontal_velocity": (("time", "space"), self.u),
+                "vertical_velocity": (("time", "space"), self.w),
+                "surface_slope": (("time", "space"), self.slope),
             },
             coords={"time": self.time, "space": self.x},
         )

From a62d246b5898380d4d488c788df248c969084858 Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Thu, 27 Feb 2025 14:56:57 -0500
Subject: [PATCH 2/7] Update README

---
 README.md | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index b53c9cc..a8c203b 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,7 @@ riding on longer waves.
 ## Features
 
 * Solves the full wave crest and action balance equations in 1-d.
-* Linear (1st order) or Stokes (3rd order) long waves
+* Wave types: 1st order linear wave, 3rd order Stokes wave, fully nonlinear wave (SSGW)
 * Infinite long-wave trains or long-wave groups
 * Effective gravity, propagation, and advection in curvilinear coordinates
 * Optionally, output all tendencies at all time steps
@@ -17,6 +17,7 @@ riding on longer waves.
 ### Install 2wave
 
 ```
+pip install git+https://github.com/wavesgroup/ssgw # dependency
 pip install git+https://github.com/wavesgroup/2wave
 ```
 
@@ -38,6 +39,7 @@ cd 2wave
 python3 -m venv venv
 source venv/bin/activate
 pip install -U .
+pip install -U https://github.com/wavesgroup/ssgw
 pytest
 ```
 

From d3120ea05828f3a8995d508bb14d3cc8af7de02c Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Fri, 28 Feb 2025 14:53:51 -0500
Subject: [PATCH 3/7] Add a Lagrangian correction to effective gravities

---
 test_twowave.py |  14 +++-
 twowave.py      | 205 +++++-------------------------------------------
 2 files changed, 31 insertions(+), 188 deletions(-)

diff --git a/test_twowave.py b/test_twowave.py
index 5a0640b..f0e4a4e 100644
--- a/test_twowave.py
+++ b/test_twowave.py
@@ -43,8 +43,18 @@ def test_elevation():
 
 
 def test_gravity():
-    assert gravity(0, 0, 0, 1, G0, wave_type="linear") == G0
-    assert gravity(0, 0, 0, 1, G0, wave_type="stokes") == G0
+    # Test with array input since gravity no longer works with scalar x
+    phase = np.linspace(0, 2 * np.pi, 128, endpoint=False)
+    k = 1
+    a = 0.1
+    x = phase / k
+    omega = np.sqrt(G0 * k)
+    assert np.isclose(
+        gravity(x, 0, a, k, omega, G0, wave_type="linear")[0], 8.83, rtol=1e-2
+    )
+    assert np.isclose(
+        gravity(x, 0, a, k, omega, G0, wave_type="stokes")[0], 8.83, rtol=1e-2
+    )
 
 
 def test_wave_modulation_model():
diff --git a/twowave.py b/twowave.py
index 0b4fd35..9c01f1e 100644
--- a/twowave.py
+++ b/twowave.py
@@ -100,67 +100,6 @@ def surface_slope(
         raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
-def gravity(
-    x: float,
-    t: float,
-    a: float,
-    k: float,
-    omega: float,
-    g0: float = 9.8,
-    wave_type: str = "linear",
-    nonlinear_props: tuple = None,
-) -> float:
-    """Gravitational acceleration at the surface of a long wave.
-    Supports linear, Stokes, and nonlinear wave types.
-    """
-    phase = k * x - omega * t
-
-    if wave_type == "linear":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        return g0 * (
-            1 - a * k * np.exp(k * eta) * (np.cos(phase) - a * k * np.sin(phase) ** 2)
-        )
-    elif wave_type == "stokes":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        return g0 * (
-            1
-            - a
-            * k
-            * (
-                np.cos(phase)
-                - a
-                * k
-                * np.sin(phase)
-                * (
-                    (1 - 1 / 16 * (a * k) ** 2) * np.sin(phase)
-                    + a * k * np.sin(2 * phase)
-                    + 9 / 8 * (a * k) ** 2 * np.sin(3 * phase)
-                )
-            )
-            * np.exp(k * eta)
-        )
-    elif wave_type == "nonlinear":
-        if nonlinear_props is None:
-            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
-        positions, _, _, _, _, az = nonlinear_props
-
-        # Handle both scalar and array inputs
-        x_array = np.atleast_1d(x)
-        result = np.zeros_like(x_array, dtype=float)
-
-        # Apply for each x value
-        for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
-            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Get effective gravity at this position
-            result[i] = g0 + np.interp(x_phase * k, positions, az, period=2 * np.pi)
-
-        # Return scalar if input was scalar
-        return result[0] if np.isscalar(x) else result
-    else:
-        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
-
-
 def orbital_horizontal_velocity(
     x: float,
     z: float,
@@ -235,122 +174,7 @@ def orbital_vertical_velocity(
         raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
 
 
-def orbital_horizontal_acceleration(
-    x: float,
-    t: float,
-    a: float,
-    k: float,
-    omega: float,
-    wave_type: str = "linear",
-    nonlinear_props: tuple = None,
-) -> float:
-    phase = k * x - omega * t
-    if wave_type == "linear":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        dU_dt = (
-            a * omega**2 * np.exp(k * eta) * np.sin(phase) * (a * k * np.cos(phase) + 1)
-        )
-        return dU_dt
-    elif wave_type == "stokes":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        term1 = a * omega**2 * np.exp(k * eta) * np.sin(phase)
-        term2 = (
-            a**2
-            * omega**2
-            * k
-            * np.exp(k * eta)
-            * np.cos(phase)
-            * (
-                np.sin(phase)
-                + a * k * np.sin(2 * phase)
-                + (a * k) ** 2
-                * (9 / 8.0 * np.sin(3 * phase) - 1 / 16.0 * np.sin(phase))
-            )
-        )
-        dU_dt = term1 + term2
-        return dU_dt
-    elif wave_type == "nonlinear":
-        if nonlinear_props is None:
-            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
-        positions, _, _, _, ax, _ = nonlinear_props
-
-        # Handle both scalar and array inputs
-        x_array = np.atleast_1d(x)
-        result = np.zeros_like(x_array, dtype=float)
-
-        # Apply for each x value
-        for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
-            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Get horizontal acceleration at this position
-            result[i] = np.interp(x_phase * k, positions, ax, period=2 * np.pi)
-
-        # Return scalar if input was scalar
-        return result[0] if np.isscalar(x) else result
-    else:
-        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
-
-
-def orbital_vertical_acceleration(
-    x: float,
-    t: float,
-    a: float,
-    k: float,
-    omega: float,
-    wave_type: str = "linear",
-    nonlinear_props: tuple = None,
-) -> float:
-    phase = k * x - omega * t
-    if wave_type == "linear":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        dW_dt = (
-            a
-            * omega**2
-            * np.exp(k * eta)
-            * (a * k * np.sin(phase) ** 2 - np.cos(phase))
-        )
-        return dW_dt
-    elif wave_type == "stokes":
-        eta = elevation(x, t, a, k, omega, wave_type)
-        term1 = -a * omega**2 * np.exp(k * eta) * np.cos(phase)
-        term2 = (
-            a**2
-            * omega**2
-            * k
-            * np.exp(k * eta)
-            * np.sin(phase)
-            * (
-                np.sin(phase)
-                + a * k * np.sin(2 * phase)
-                + (a * k) ** 2
-                * (9 / 8.0 * np.sin(3 * phase) - 1 / 16.0 * np.sin(phase))
-            )
-        )
-        dW_dt = term1 + term2
-        return dW_dt
-    elif wave_type == "nonlinear":
-        if nonlinear_props is None:
-            raise ValueError("nonlinear_props must be provided for nonlinear wave type")
-        positions, _, _, _, _, az = nonlinear_props
-
-        # Handle both scalar and array inputs
-        x_array = np.atleast_1d(x)
-        result = np.zeros_like(x_array, dtype=float)
-
-        # Apply for each x value
-        for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
-            x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Get vertical acceleration at this position
-            result[i] = np.interp(x_phase * k, positions, az, period=2 * np.pi)
-
-        # Return scalar if input was scalar
-        return result[0] if np.isscalar(x) else result
-    else:
-        raise ValueError("wave_type must be 'linear', 'stokes', or 'nonlinear'")
-
-
-def gravity_curvilinear(
+def gravity(
     x: float,
     t: float,
     a: float,
@@ -359,13 +183,20 @@ def gravity_curvilinear(
     g0: float = 9.8,
     wave_type: str = "linear",
     nonlinear_props: tuple = None,
+    curvilinear: bool = True,
 ) -> float:
-    dU_dt = orbital_horizontal_acceleration(
-        x, t, a, k, omega, wave_type, nonlinear_props
-    )
-    dW_dt = orbital_vertical_acceleration(x, t, a, k, omega, wave_type, nonlinear_props)
-    slope = surface_slope(x, t, a, k, omega, wave_type, nonlinear_props)
-    g = g0 * np.cos(slope) + dW_dt * np.cos(slope) + dU_dt * np.sin(slope)
+    z = elevation(x, t, a, k, omega, wave_type, nonlinear_props)
+    U = orbital_horizontal_velocity(x, z, t, a, k, omega, wave_type, nonlinear_props)
+    W = orbital_vertical_velocity(x, z, t, a, k, omega, wave_type, nonlinear_props)
+    Cp = omega / k
+    dx = np.diff(x)[0]
+    a_z = -diff(W) / dx * (Cp - U)
+    if curvilinear:
+        a_x = -diff(U) / dx * (Cp - U)
+        slope = surface_slope(x, t, a, k, omega, wave_type, nonlinear_props)
+        g = g0 * np.cos(slope) + a_z * np.cos(slope) + a_x * np.sin(slope)
+    else:
+        g = g0 + a_z
     return g
 
 
@@ -387,8 +218,8 @@ def nonlinear_wave_properties(
 
     dx = diff(x)
     dx[dx < 0] += np.pi
-    a_z = -diff(w) / dx * Cp
-    a_x = -diff(u) / dx * Cp
+    a_z = -diff(w) / dx * (Cp - u)
+    a_x = -diff(u) / dx * (Cp - u)
 
     # Shift back by pi
     x += np.pi
@@ -518,7 +349,7 @@ def run(
             )
 
         self.elevation = elevation
-        self.gravity = gravity_curvilinear if self.curvilinear else gravity
+        self.gravity = gravity
         self._wave_type = wave_type
 
         if wave_type == "nonlinear":
@@ -671,6 +502,7 @@ def wavenumber_tendency(self, k, t):
             self.grav0,
             wave_type=self._wave_type,
             nonlinear_props=self.nonlinear_props,
+            curvilinear=self.curvilinear,
         )
         self.g[self.current_time_step] = g
 
@@ -771,6 +603,7 @@ def waveaction_tendency(self, N, t):
             self.grav0,
             wave_type=self._wave_type,
             nonlinear_props=self.nonlinear_props,
+            curvilinear=self.curvilinear,
         )
         Cg = (
             angular_frequency(g, self.k[self.current_time_step])

From 88fb1316a81d5e0be3e4a76f7af2ca19c7afac3f Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Mon, 3 Mar 2025 13:21:25 -0500
Subject: [PATCH 4/7] Install SSGW as a dependency

---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 177c150..35eddeb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,9 @@ authors = [
 dependencies = [
     "numpy",
     "rich",
+    "scipy",
     "xarray",
+    "ssgw @ git+https://github.com/wavesgroup/ssgw",
 ]
 
 [tool.setuptools]

From bda000c860d6aaf9a309067731ca6fc0e46bee35 Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Tue, 4 Mar 2025 12:07:36 -0500
Subject: [PATCH 5/7] Slight optimization in nonlinear mode

---
 twowave.py | 24 +++++++-----------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/twowave.py b/twowave.py
index 9c01f1e..eaa6eaf 100644
--- a/twowave.py
+++ b/twowave.py
@@ -39,11 +39,9 @@ def elevation(
         x_array = np.atleast_1d(x)
         result = np.zeros_like(x_array, dtype=float)
 
-        # Apply for each x value
+        # Apply a phase shift of the elevations by interpolating
         for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
             x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Interpolate elevation at the phase-shifted position
             result[i] = np.interp(x_phase * k, positions, elevations, period=2 * np.pi)
 
         # Return scalar if input was scalar
@@ -82,17 +80,13 @@ def surface_slope(
         result = np.zeros_like(x_array, dtype=float)
 
         # Calculate dense slopes once
-        dense_pos = np.linspace(0, 2 * np.pi, 1000)
-        dense_elev = np.interp(dense_pos, positions, elevations, period=2 * np.pi)
-        dx = dense_pos[1] - dense_pos[0]
-        slopes = np.gradient(dense_elev, dx)
+        dx = positions[1] - positions[0]
+        slopes = np.gradient(elevations, dx)
 
-        # Apply for each x value
+        # Apply a phase shift of the slopes by interpolating
         for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
             x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Interpolate slope at the requested position
-            result[i] = np.interp(x_phase * k, dense_pos, slopes, period=2 * np.pi) / k
+            result[i] = np.interp(x_phase * k, positions, slopes, period=2 * np.pi) / k
 
         # Return scalar if input was scalar
         return result[0] if np.isscalar(x) else result
@@ -124,11 +118,9 @@ def orbital_horizontal_velocity(
         x_array = np.atleast_1d(x)
         result = np.zeros_like(x_array, dtype=float)
 
-        # Apply for each x value
+        # Apply a phase shift of the horizontal velocities by interpolating
         for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
             x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Get horizontal velocity at this position
             result[i] = np.interp(x_phase * k, positions, u, period=2 * np.pi)
 
         # Return scalar if input was scalar
@@ -161,11 +153,9 @@ def orbital_vertical_velocity(
         x_array = np.atleast_1d(x)
         result = np.zeros_like(x_array, dtype=float)
 
-        # Apply for each x value
+        # Apply a phase shift of the vertical velocities by interpolating
         for i, xi in enumerate(x_array):
-            # Apply phase shift based on time
             x_phase = (xi - omega / k * t) % (2 * np.pi / k)
-            # Get vertical velocity at this position
             result[i] = np.interp(x_phase * k, positions, w, period=2 * np.pi)
 
         # Return scalar if input was scalar

From c3a1135fd8921fe512d8890d7e12453f134821e6 Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Tue, 4 Mar 2025 12:47:10 -0500
Subject: [PATCH 6/7] Take into account group velocity in the Lagrangian
 evaluation of effective gravity

---
 twowave.py | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/twowave.py b/twowave.py
index eaa6eaf..5f7cb81 100644
--- a/twowave.py
+++ b/twowave.py
@@ -171,6 +171,7 @@ def gravity(
     k: float,
     omega: float,
     g0: float = 9.8,
+    cg: float = 0.0,
     wave_type: str = "linear",
     nonlinear_props: tuple = None,
     curvilinear: bool = True,
@@ -180,9 +181,9 @@ def gravity(
     W = orbital_vertical_velocity(x, z, t, a, k, omega, wave_type, nonlinear_props)
     Cp = omega / k
     dx = np.diff(x)[0]
-    a_z = -diff(W) / dx * (Cp - U)
+    a_z = -diff(W) / dx * (Cp - U - cg)
     if curvilinear:
-        a_x = -diff(U) / dx * (Cp - U)
+        a_x = -diff(U) / dx * (Cp - U - cg)
         slope = surface_slope(x, t, a, k, omega, wave_type, nonlinear_props)
         g = g0 * np.cos(slope) + a_z * np.cos(slope) + a_x * np.sin(slope)
     else:
@@ -360,9 +361,11 @@ def run(
         self.k = np.zeros((self.num_time_steps, self.grid_size), dtype=np.float32)
         self.a = np.zeros_like(self.k)
         self.N = np.zeros_like(self.k)
+        self.cg = np.zeros_like(self.k)
         self.k[0] = self.k_short
         self.a[0] = self.a_short
         self.N[0] = 1  # FIXME: This is a placeholder value.
+        self.cg[0] = 0.5 * angular_frequency(self.grav0, self.k[0]) / self.k[0]
 
         # Allocate and initialize short-wave diagnostic fields.
         self.g = np.zeros_like(self.k)
@@ -490,6 +493,7 @@ def wavenumber_tendency(self, k, t):
             self.k_long,
             self.omega_long,
             self.grav0,
+            self.cg[self.current_time_step],
             wave_type=self._wave_type,
             nonlinear_props=self.nonlinear_props,
             curvilinear=self.curvilinear,
@@ -500,6 +504,8 @@ def wavenumber_tendency(self, k, t):
         self.omega[self.current_time_step] = omega
 
         Cg = omega / k / 2
+        self.cg[self.current_time_step] = Cg
+
         k_propagation_tendency = -Cg * diff(k) / self.ds
         k_advection_tendency = -vel * diff(k) / self.ds
         k_convergence_tendency = -k * diff(vel) / self.ds
@@ -523,11 +529,6 @@ def waveaction_tendency(self, N, t):
         """Compute the tendencies of the wave action balance at time t."""
         eta_ramp = self.get_elevation_ramp(t)
 
-        if self._wave_type == "nonlinear":
-            self.nonlinear_props = nonlinear_wave_properties(
-                eta_ramp * self.a_long, self.k_long, self.grav0
-            )
-
         eta = self.elevation(
             self.x,
             t,
@@ -591,15 +592,19 @@ def waveaction_tendency(self, N, t):
             self.k_long,
             self.omega_long,
             self.grav0,
+            self.cg[self.current_time_step],
             wave_type=self._wave_type,
             nonlinear_props=self.nonlinear_props,
             curvilinear=self.curvilinear,
         )
+
         Cg = (
             angular_frequency(g, self.k[self.current_time_step])
             / self.k[self.current_time_step]
             / 2
         )
+        self.cg[self.current_time_step] = Cg
+
         N_propagation_tendency = -Cg * diff(N) / self.ds
         N_advection_tendency = -vel * diff(N) / self.ds
         N_convergence_tendency = -N * diff(vel) / self.ds

From 8ab8a5443573d1e92ab4bec2889409808e6a588c Mon Sep 17 00:00:00 2001
From: milancurcic <caomaco@gmail.com>
Date: Tue, 4 Mar 2025 12:51:20 -0500
Subject: [PATCH 7/7] Update copyright year and README

---
 LICENSE   |  2 +-
 README.md | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/LICENSE b/LICENSE
index 278383c..b70d3e0 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 Milan Curcic
+Copyright (c) 2024-2025 Milan Curcic
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/README.md b/README.md
index a8c203b..017e419 100644
--- a/README.md
+++ b/README.md
@@ -5,19 +5,21 @@ riding on longer waves.
 
 ## Features
 
-* Solves the full wave crest and action balance equations in 1-d.
-* Wave types: 1st order linear wave, 3rd order Stokes wave, fully nonlinear wave (SSGW)
+* Solves the nonlinear wave crest and action balance equations in 1-d.
+* Wave types:
+  - 1st order linear wave
+  - 3rd order Stokes wave
+  - Fully nonlinear wave using [SSGW](https://github.com/wavesgroup/ssgw) by Clamond & Dutykh (2018, JFM)
 * Infinite long-wave trains or long-wave groups
 * Effective gravity, propagation, and advection in curvilinear coordinates
-* Optionally, output all tendencies at all time steps
 * Output as Xarray Dataset
+* Optionally, output all tendencies at all time steps
 
 ## Getting started
 
 ### Install 2wave
 
 ```
-pip install git+https://github.com/wavesgroup/ssgw # dependency
 pip install git+https://github.com/wavesgroup/2wave
 ```
 
@@ -39,7 +41,6 @@ cd 2wave
 python3 -m venv venv
 source venv/bin/activate
 pip install -U .
-pip install -U https://github.com/wavesgroup/ssgw
 pytest
 ```