From baa00a4c371a501d1f5f27afbdbd8f9948b4dee4 Mon Sep 17 00:00:00 2001
From: Gavin Macaulay <gavin@macaulay.co.nz>
Date: Fri, 29 Nov 2024 11:54:40 +1300
Subject: [PATCH] Add utility functions to make shapes for the dwba model

---
 docs/api_reference.md    |  2 +
 src/echosms/__init__.py  | 11 ++++--
 src/echosms/utilsdwba.py | 82 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 4 deletions(-)
 create mode 100644 src/echosms/utilsdwba.py

diff --git a/docs/api_reference.md b/docs/api_reference.md
index d0f3688..eaf1f5d 100644
--- a/docs/api_reference.md
+++ b/docs/api_reference.md
@@ -61,3 +61,5 @@ There are several models that use the distorted-wave Born approximation, documen
 ## Utilities
 
 ::: echosms.utils
+
+::: echosms.utilsdwba
diff --git a/src/echosms/__init__.py b/src/echosms/__init__.py
index f90f35a..92d7d0a 100644
--- a/src/echosms/__init__.py
+++ b/src/echosms/__init__.py
@@ -1,7 +1,8 @@
 """Setup the public API for echoSMs."""
-from .utils import wavenumber, wavelength, Neumann, h1, prolate_swf, spherical_jnpp, split_dict, theoretical_Sa
+from .utils import wavenumber, wavelength, Neumann, h1, prolate_swf, spherical_jnpp
 from .utils import pro_rad1, pro_rad2, pro_ang1
-from .utils import as_dataframe, as_dataarray, as_dict
+from .utils import as_dataframe, as_dataarray, as_dict, split_dict, theoretical_Sa
+from .utilsdwba import create_dwba_spheroid, create_dwba_cylinder
 from .scattermodelbase import ScatterModelBase
 from .benchmarkdata import BenchmarkData
 from .referencemodels import ReferenceModels
@@ -17,6 +18,8 @@
 __all__ = ['ScatterModelBase', 'BenchmarkData', 'ReferenceModels',
            'MSSModel', 'PSMSModel', 'DCMModel', 'ESModel', 'PTDWBAModel',
            'DWBAModel', 'SDWBAModel', 'KAModel',
-           'wavenumber', 'wavelength', 'Neumann', 'h1', 'spherical_jnpp', 'prolate_swf', 'theoretical_Sa',
+           'wavenumber', 'wavelength', 'Neumann', 'h1', 'spherical_jnpp', 'prolate_swf',
+           'theoretical_Sa',
            'pro_rad1', 'pro_rad2', 'pro_ang1',
-           'as_dataframe', 'as_dataarray', 'as_dict', 'split_dict']
+           'as_dataframe', 'as_dataarray', 'as_dict', 'split_dict',
+           'create_dwba_spheroid', 'create_dwba_cylinder']
diff --git a/src/echosms/utilsdwba.py b/src/echosms/utilsdwba.py
new file mode 100644
index 0000000..8ef846d
--- /dev/null
+++ b/src/echosms/utilsdwba.py
@@ -0,0 +1,82 @@
+"""Miscellaneous functions for the DWBA models."""
+import numpy as np
+
+
+def create_dwba_spheroid(major_radius: float, minor_radius: float, spacing: float = 0.0001):
+    """Create shape description variables for the DWBA model for spheres and prolate spheroids.
+
+    The shape descriptions are essentially a set of discs and their orientation.
+
+    Notes
+    -----
+    Currently only supports prolate spheroids and spheres (set `major_radius` and `minor_radius`
+    to the same value to get a sphere).
+
+    Parameters
+    ----------
+    major_radius :
+        The major radius [m] of the spheroid.
+    minor_radius :
+        The minor radius [m] of the spheroid.
+    spacing :
+        The spacing [m] between successive discs.
+
+    Returns
+    -------
+    rv_pos : iterable[np.ndarray]
+        An interable of vectors of the 3D positions of the centre of each disc that
+        defines the spheroid. Each vector has three values corresponding to
+        the _x_, _y_, and _z_ coordinates [m] of the disc centre.
+    rv_tan : iterable[np.ndarray]
+        An interable of unit vectors of the tangent to the target body axis at
+        the points given in `rv_pos`. Each vector has three values corresponding to
+        the _x_, _y_, and _z_ components of the tangent vector.
+    a : iterable
+        The radii [m] of the discs that define the spheroid.
+    """
+    v = np.linspace(0, np.pi, int(round(2*major_radius/spacing)))
+    a = minor_radius*np.sin(v)  # radius at points along the spheroid
+    x = major_radius-major_radius*np.cos(v)  # shift so that origin is at one end
+
+    # List of disc position vectors
+    rv_pos = [np.array([i, 0, 0]) for i in x]
+
+    # List of tangent vectors to sphere axis (all the same for spheroids)
+    rv_tan = [np.array([1, 0, 0])] * len(x)
+
+    return rv_pos, rv_tan, a
+
+
+def create_dwba_cylinder(radius: float, length: float, spacing: float = 0.0001):
+    """Create shape description variables for the DWBA model for cylinders.
+
+    The shape descriptions are essentially a set of discs and their orientation.
+
+    Parameters
+    ----------
+    radius :
+        The radius [m] of the cylinder.
+    length :
+        The length [m] of the cylinder.
+    spacing :
+        The spacing [m] between successive discs.
+
+    Returns
+    -------
+    rv_pos : iterable[np.ndarray]
+        An interable of vectors of the 3D positions of the centre of each disc that
+        defines the cylinder. Each vector has three values corresponding to
+        the _x_, _y_, and _z_ coordinates [m] of the disc centre.
+    rv_tan : iterable[np.ndarray]
+        An interable of unit vectors of the tangent to the cylinder body axis at
+        the points given in `rv_pos`. Each vector has three values corresponding to
+        the _x_, _y_, and _z_ components of the tangent vector.
+    a : iterable
+        The radii [m] of the discs that define the spheroid.
+    """
+    pos = np.linspace(0, length, int(round(length/spacing)))
+    rv_pos = [np.array([x, 0, 0]) for x in pos]
+    rv_tan = [np.array([1, 0, 0])] * len(pos)
+    a = [radius] * len(pos)
+
+    return rv_pos, rv_tan, a