From 78b5aa131ee3d02bd1783272b57578fdf980405c Mon Sep 17 00:00:00 2001
From: Daniel da Silva <Daniel.e.daSilva@nasa.gov>
Date: Sat, 23 Nov 2024 02:14:41 -0500
Subject: [PATCH] Drop dependency on ai.cs and fork functions used there
 (theyre simple)

---
 pyproject.toml                |  1 -
 rbinvariantslib/invariants.py |  3 +-
 rbinvariantslib/models.py     |  7 ----
 rbinvariantslib/utils.py      | 64 +++++++++++++++++++++++++++++++++--
 4 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8cf808d..6263b7a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,7 +24,6 @@ classifiers=[
 ]
 
 dependencies = [
-    "ai.cs>=1.0.7",
     "astropy>=5.1",
     "cdasws>=1.8.7",
     "h5py>=3.10.0",
diff --git a/rbinvariantslib/invariants.py b/rbinvariantslib/invariants.py
index f0a81ec..bdafbd3 100644
--- a/rbinvariantslib/invariants.py
+++ b/rbinvariantslib/invariants.py
@@ -2,7 +2,6 @@
 from dataclasses import dataclass
 from typing import Any, List, Optional, Tuple
 
-from ai import cs
 import numpy as np
 from numpy.typing import NDArray
 from scipy import interpolate
@@ -233,7 +232,7 @@ def calculate_K(
 
     # Get the trace latitudes and Bm if not specified
     # ------------------------------------------------------------------------
-    _, trace_latitude, _ = cs.cart2sp(
+    _, trace_latitude, _ = utils.cart2sp(
         x=trace.points[:, 0], y=trace.points[:, 1], z=trace.points[:, 2]
     )
     trace_sorter = np.argsort(trace_latitude)
diff --git a/rbinvariantslib/models.py b/rbinvariantslib/models.py
index 24cb17f..6f08455 100644
--- a/rbinvariantslib/models.py
+++ b/rbinvariantslib/models.py
@@ -11,8 +11,6 @@
 import os
 from typing import cast, Dict, List, Tuple, Union
 
-
-from ai import cs
 from astropy import constants, units
 from cdasws import CdasWs
 import h5py
@@ -102,11 +100,6 @@ def __init__(self, x, y, z, Bx, By, Bz, inner_boundary):
         self._mesh = pv.StructuredGrid(x, y, z)
         self._mesh.point_data["B"] = B
 
-        #R_grid, theta_grid, phi_grid = cs.cart2sp(x, y, z)
-        #self._mesh.point_data["R_grid"] = R_grid.flatten(order="F")
-        #self._mesh.point_data["Theta_grid"] = theta_grid.flatten(order="F")
-        #self._mesh.point_data["Phi_grid"] = phi_grid.flatten(order="F")
-
     def trace_field_line(
         self, starting_point, step_size
     ) -> FieldLineTrace:
diff --git a/rbinvariantslib/utils.py b/rbinvariantslib/utils.py
index ed2bc8e..d19d5ce 100644
--- a/rbinvariantslib/utils.py
+++ b/rbinvariantslib/utils.py
@@ -34,7 +34,7 @@ def sp2cart_point(r: float, phi: float, theta: float) -> Tuple[float, float, flo
     Returns
       x, y, z: Cartesian coordinates
     """
-    point = cs.sp2cart(r=r, phi=phi, theta=theta)  # returns tuple of 0d arrays
+    point = sp2cart(r=r, phi=phi, theta=theta)  # returns tuple of 0d arrays
     point = tuple(np.array(point).tolist())
 
     return point
@@ -50,12 +50,72 @@ def cart2sp_point(x: float, y: float, z: float) -> Tuple[float, float, float]:
       phi: longitude
       theta: latitude
     """
-    point = cs.cart2sp(x=x, y=y, z=z)  # returns tuple of 0d arrays
+    point = cart2sp(x=x, y=y, z=z)  # returns tuple of 0d arrays
     point = tuple(np.array(point).tolist())
 
     return point
 
 
+# Function forked from https://bitbucket.org/isavnin/ai.cs/src/master/src/ai/cs.py
+# because broken in Python 3.12
+def cart2sp(x, y, z):
+    """Converts data from cartesian coordinates into spherical.
+
+    Args:
+        x (scalar or array_like): X-component of data.
+        y (scalar or array_like): Y-component of data.
+        z (scalar or array_like): Z-component of data.
+
+    Returns:
+        Tuple (r, theta, phi) of data in spherical coordinates.
+    """
+    x = np.asarray(x)
+    y = np.asarray(y)
+    z = np.asarray(z)
+    scalar_input = False
+    if x.ndim == 0 and y.ndim == 0 and z.ndim == 0:
+        x = x[None]
+        y = y[None]
+        z = z[None]
+        scalar_input = True
+    r = np.sqrt(x**2+y**2+z**2)
+    theta = np.arcsin(z/r)
+    phi = np.arctan2(y, x)
+    if scalar_input:
+        return (r.squeeze(), theta.squeeze(), phi.squeeze())
+    return (r, theta, phi)
+
+
+# Function forked from https://bitbucket.org/isavnin/ai.cs/src/master/src/ai/cs.py
+# because broken in Python 3.12
+def sp2cart(r, theta, phi):
+    """Converts data in spherical coordinates into cartesian.
+
+    Args:
+        r (scalar or array_like): R-component of data.
+        theta (scalar or array_like): Theta-component of data.
+        phi (scalar or array_like): Phi-component of data.
+
+    Returns:
+        Tuple (x, y, z) of data in cartesian coordinates.
+    """
+    r = np.asarray(r)
+    theta = np.asarray(theta)
+    phi = np.asarray(phi)
+    scalar_input = False
+    if r.ndim == 0 and theta.ndim == 0 and phi.ndim == 0:
+        r = r[None]
+        theta = theta[None]
+        phi = phi[None]
+        scalar_input = True
+    x = r*np.cos(theta)*np.cos(phi)
+    y = r*np.cos(theta)*np.sin(phi)
+    z = r*np.sin(theta)
+    if scalar_input:
+        return (x.squeeze(), y.squeeze(), z.squeeze())
+    return (x, y, z)
+
+
 def lfm_get_eq_slice(data):
     """Gets an equitorial slice from data on an LFM grid.