coord_system3d: simplify

Author: paugier

fluidsim/operators/coord_system3d.py

@@ -20,60 +20,38 @@ def __init__(self, x, y, z):
         self.y = y
         self.z = z
 
+        # tiny constant
+        EPSILON = 1e-12
+
         # cylindrical coordinates
         self.rh = np.sqrt(self.x**2 + self.y**2)
-        # rtheta is defined between [-pi; pi] and is 0 for x = y = 0
-        self.r_theta = np.atan2(self.y, self.x)
+        self.rh_not0 = np.where(self.rh != 0, self.rh, EPSILON)
 
         # spherical coordinates
-        self.r = np.sqrt(self.x**2 + self.y**2 + self.z**2)
+        r = np.sqrt(self.x**2 + self.y**2 + self.z**2)
+        self.r_not0 = np.where(r != 0, r, EPSILON)
 
-    def compute_cylindrical_components(self, vx, vy, vz):
-        """
-        Convert (vx, vy, vz) in cylindrical coordinates (vh, vt, vz)
-        """
+    def compute_r_theta(self):
+        """r_theta is defined between [-pi; pi] and is 0 for x = y = 0"""
+        return np.atan2(self.y, self.x)
 
-        with np.errstate(divide="ignore", invalid="ignore"):
-            vh = np.where(self.rh != 0, (self.x * vx + self.y * vy) / self.rh, 0)
-            vt = np.where(self.rh != 0, (-self.y * vx + self.x * vy) / self.rh, 0)
+    def compute_cylindrical_components(self, vx, vy, vz):
+        """Convert (vx, vy, vz) in cylindrical coordinates (vh, vt, vz)"""
 
+        vh = (self.x * vx + self.y * vy) / self.rh_not0
+        vt = (-self.y * vx + self.x * vy) / self.rh_not0
         return vh, vt, vz
 
     def compute_radial_component(self, vx, vy, vz):
-        """
-        Compute radial component vr along r
-        """
-
-        with np.errstate(divide="ignore", invalid="ignore"):
-            vr = np.where(
-                self.r != 0, (self.x * vx + self.y * vy + self.z * vz) / self.r, 0
-            )
-
-        return vr
+        """Compute radial component"""
+        return (self.x * vx + self.y * vy + self.z * vz) / self.r_not0
 
     def compute_spherical_components(self, vx, vy, vz):
-        """
-        Convert vx, vy, vz in spherical coordinates (vr, vt, vp)
-        """
-
-        with np.errstate(divide="ignore", invalid="ignore"):
-            self.rp = np.where(
-                self.r != 0, np.arccos(np.clip(self.z / self.r, -1, 1)), 0
-            )
-
-            vr = np.where(
-                self.r != 0, (self.x * vx + self.y * vy + self.z * vz) / self.r, 0
-            )
-
-            vt = np.where(self.rh != 0, (-self.y * vx + self.x * vy) / self.rh, 0)
-
-            vp = np.where(
-                self.r != 0,
-                (self.z * (self.x * vx + self.y * vy) - self.rh**2 * vz)
-                / (self.r * self.rh),
-                0,
-            )
-            # Alternative for vp when rh = 0
-            vp = np.where(self.rh != 0, vp, 0)
-
+        """Convert (vx, vy, vz) in spherical coordinates (vr, vt, vp)"""
+        vr = self.compute_radial_component(vx, vy, vz)
+        vt = (-self.y * vx + self.x * vy) / self.rh_not0
+        vp = (
+            (self.z * (self.x * vx + self.y * vy) - self.rh**2 * vz)
+            / (self.r_not0 * self.rh_not0),
+        )
         return vr, vt, vp

fluidsim/operators/test/test_coord_system3d.py

@@ -48,3 +48,9 @@ def test_coord_system_converter(vector_kind, converter, allclose):
             pass
         case _:
             raise ValueError
+
+
+def test_compute_r_theta(converter, allclose):
+    r_theta = converter.compute_r_theta()
+    # TODO: fix this
+    assert allclose(r_theta, r_theta)