more centring tweaks; add import test

.github/workflows/continuousintegration.yaml

@@ -22,6 +22,8 @@ jobs:
         run: |  
           python -m pip install --upgrade pip  
           pip install -r test_requirements.txt
+        run: |
+          python tests/importtest.py
 #      - name: Test with pytest  
 #        run: |  
 #         coverage run -m pytest  -v -s  

exptool/analysis/__init__.py

@@ -1,7 +1,14 @@
 
 
-# this is EXPtool
 
-__all__ = ['commensurability','directdetection','ellipsetools','forcetrace','globalquantities','pattern','trapping']
+__all__ = ['centering',
+           'commensurability',
+           'directdetection',
+           'ellipsetools',
+           'globalquantities',
+           'pattern',
+           'resonancemodel',
+           'rotationcurves',
+           'trapping']
 
 

exptool/analysis/centering.py

@@ -13,6 +13,7 @@
     
 09 Jul 2021: First introduction
 18 Sep 2024: Class structure overhaul
+23 Oct 2024: Added radius masking to alignment
 """
 import numpy as np
 
@@ -85,7 +86,7 @@ def __init__(self, x, y, z, vx, vy, vz, mass):
         self.vz = vz
         self.mass = mass
 
-    def recentre_pos_and_vel(self, r_max):
+    def recentre_pos_and_vel(self, r_max=np.inf):
         """
         Recenter the position and velocity of the particles based on the center of mass 
         and mass-weighted velocity within r_max.
@@ -99,84 +100,22 @@ def recentre_pos_and_vel(self, r_max):
         """
         mask = (self.x**2 + self.y**2 + self.z**2) < r_max**2
         mass_tot = np.sum(self.mass[mask])
+
         x_cm = np.sum(self.x[mask]*self.mass[mask])/mass_tot
         y_cm = np.sum(self.y[mask]*self.mass[mask])/mass_tot
         z_cm = np.sum(self.z[mask]*self.mass[mask])/mass_tot
+
         vx_cm = np.sum(self.vx[mask]*self.mass[mask])/mass_tot
         vy_cm = np.sum(self.vy[mask]*self.mass[mask])/mass_tot
         vz_cm = np.sum(self.vz[mask]*self.mass[mask])/mass_tot
-        return self.x - x_cm, self.y - y_cm, self.z - z_cm, self.vx - vx_cm, self.vy - vy_cm, self.vz - vz_cm
-
-    def rotate(self, axis, angle):
-        """
-        Rotate the particles around a given axis by a specified angle.
-
-        Parameters:
-        axis : array
-            Axis of rotation.
-        angle : float
-            Angle of rotation in radians.
-
-        Returns:
-        Rotated positions.
-        """
-        axisx, axisy, axisz = axis
-
-        # Rotation for position vector (x, y, z)
-        dot_pos = self.x * axisx + self.y * axisy + self.z * axisz  # Dot product for position
-        crossx_pos = axisy * self.z - axisz * self.y  # Cross product for position
-        crossy_pos = axisz * self.x - axisx * self.z
-        crossz_pos = axisx * self.y - axisy * self.x
-
-        cosa = np.cos(angle)
-        sina = np.sin(angle)
-
-        x_rot = self.x * cosa + crossx_pos * sina + axisx * dot_pos * (1 - cosa)
-        y_rot = self.y * cosa + crossy_pos * sina + axisy * dot_pos * (1 - cosa)
-        z_rot = self.z * cosa + crossz_pos * sina + axisz * dot_pos * (1 - cosa)
-
-        # Rotation for velocity vector (vx, vy, vz)
-        dot_vel = self.vx * axisx + self.vy * axisy + self.vz * axisz  # Dot product for velocity
-        crossx_vel = axisy * self.vz - axisz * self.vy  # Cross product for velocity
-        crossy_vel = axisz * self.vx - axisx * self.vz
-        crossz_vel = axisx * self.vy - axisy * self.vx
-
-        vx_rot = self.vx * cosa + crossx_vel * sina + axisx * dot_vel * (1 - cosa)
-        vy_rot = self.vy * cosa + crossy_vel * sina + axisy * dot_vel * (1 - cosa)
-        vz_rot = self.vz * cosa + crossz_vel * sina + axisz * dot_vel * (1 - cosa)
-
-        return x_rot, y_rot, z_rot, vx_rot, vy_rot, vz_rot
-
-    def compute_rotation_to_vec(self, vec):
-        """
-        Compute the rotation axis and angle to align the angular momentum of the system 
-        with a given vector.
 
-        Parameters:
-        vec : array
-            Target vector to align with.
-
-        Returns:
-        axis : array
-            Rotation axis.
-        angle : float
-            Angle to rotate.
-        """
-        Lxtot = np.sum(self.y * self.vz * self.mass - self.z * self.vy * self.mass)
-        Lytot = np.sum(self.z * self.vx * self.mass - self.x * self.vz * self.mass)
-        Lztot = np.sum(self.x * self.vy * self.mass - self.y * self.vx * self.mass)
-
-        L = np.array([Lxtot, Lytot, Lztot])
-        L /= np.linalg.norm(L)
-        vec /= np.linalg.norm(vec)
-
-        axis = np.cross(L, vec)
-        axis /= np.linalg.norm(axis)
-
-        c = np.dot(L, vec)
-        angle = np.arccos(np.clip(c, -1, 1))
-
-        return axis, angle
+        # apply changes in place
+        self.x -= x_cm
+        self.y -= y_cm
+        self.z -= z_cm
+        self.vx -= vx_cm
+        self.vy -= vy_cm
+        self.vz -= vz_cm
 
     def shrinking_sphere(self, w, rmin=1., stepsize=0.5, tol=0.001, verbose=0):
         """
@@ -212,7 +151,9 @@ def shrinking_sphere(self, w, rmin=1., stepsize=0.5, tol=0.001, verbose=0):
         if verbose:
             print(f'Initial guess: {tshiftx:.0f}, {tshifty:.0f}, {tshiftz:.0f}')
 
+        nstep = 0
         while rmax > rmin:
+            nstep += 1
             u = np.where(rval < stepsize * rmax)[0]
             if float(u.size) / float(self.x.size) < tol:
                 print(f'Too few particles to continue at radius ratio {stepsize * rmax / rmax0}')
@@ -233,6 +174,9 @@ def shrinking_sphere(self, w, rmin=1., stepsize=0.5, tol=0.001, verbose=0):
             rval = np.sqrt(self.x**2 + self.y**2 + self.z**2)
             rmax *= stepsize
 
+        if nstep == 0:
+            print('No shrinking done. Check that rmin is larger than the minimum particle radius.')
+
         comvx = np.nanmedian(np.nansum(w[u] * self.vx[u]) / np.nansum(w[u]))
         comvy = np.nanmedian(np.nansum(w[u] * self.vy[u]) / np.nansum(w[u]))
         comvz = np.nanmedian(np.nansum(w[u] * self.vz[u]) / np.nansum(w[u]))
@@ -245,7 +189,80 @@ def shrinking_sphere(self, w, rmin=1., stepsize=0.5, tol=0.001, verbose=0):
         self.vy -= comvy
         self.vz -= comvz
 
-        return self.x, self.y, self.z, self.vx, self.vy, self.vz
+
+    def compute_rotation_to_vec(self, vec, rmax=np.inf):
+        """
+        Compute the rotation axis and angle to align the angular momentum of the system 
+        with a given vector.
+
+        Parameters:
+        vec : array
+            Target vector to align with.
+
+        Returns:
+        axis : array
+            Rotation axis.
+        angle : float
+            Angle to rotate.
+        """
+        mask = (self.x**2 + self.y**2 + self.z**2) < r_max**2
+
+        Lxtot = np.sum((self.mass * (self.y * self.vz - self.z * self.vy))[mask])
+        Lytot = np.sum((self.mass * (self.z * self.vx - self.x * self.vz))[mask])
+        Lztot = np.sum((self.mass * (self.x * self.vy - self.y * self.vx))[mask])
+
+        L = np.array([Lxtot, Lytot, Lztot])
+        L /= np.linalg.norm(L)
+
+        vec /= np.linalg.norm(vec)
+
+        axis = np.cross(L, vec)
+        axis /= np.linalg.norm(axis)
+
+        c = np.dot(L, vec)
+        angle = np.arccos(np.clip(c, -1, 1))
+
+        return axis, angle
+
+    def rotate(self, axis, angle):
+        """
+        Rotate the particles around a given axis by a specified angle.
+
+        Parameters:
+        axis : array
+            Axis of rotation.
+        angle : float
+            Angle of rotation in radians.
+
+        Returns:
+        Rotated positions.
+        """
+        axisx, axisy, axisz = axis
+
+        # Rotation for position vector (x, y, z)
+        dot_pos = self.x * axisx + self.y * axisy + self.z * axisz  # Dot product for position
+        crossx_pos = axisy * self.z - axisz * self.y  # Cross product for position
+        crossy_pos = axisz * self.x - axisx * self.z
+        crossz_pos = axisx * self.y - axisy * self.x
+
+        cosa = np.cos(angle)
+        sina = np.sin(angle)
+
+        x_rot = self.x * cosa + crossx_pos * sina + axisx * dot_pos * (1 - cosa)
+        y_rot = self.y * cosa + crossy_pos * sina + axisy * dot_pos * (1 - cosa)
+        z_rot = self.z * cosa + crossz_pos * sina + axisz * dot_pos * (1 - cosa)
+
+        # Rotation for velocity vector (vx, vy, vz)
+        dot_vel = self.vx * axisx + self.vy * axisy + self.vz * axisz  # Dot product for velocity
+        crossx_vel = axisy * self.vz - axisz * self.vy  # Cross product for velocity
+        crossy_vel = axisz * self.vx - axisx * self.vz
+        crossz_vel = axisx * self.vy - axisy * self.vx
+
+        vx_rot = self.vx * cosa + crossx_vel * sina + axisx * dot_vel * (1 - cosa)
+        vy_rot = self.vy * cosa + crossy_vel * sina + axisy * dot_vel * (1 - cosa)
+        vz_rot = self.vz * cosa + crossz_vel * sina + axisz * dot_vel * (1 - cosa)
+
+        return x_rot, y_rot, z_rot, vx_rot, vy_rot, vz_rot
 
     def compute_density_profile(self, R, W, rbins=10.**np.linspace(-3.7, 0.3, 100)):
         """

exptool/analysis/resonancemodel.py

@@ -59,15 +59,15 @@ def denom(r,E,J,model):
 
 def make_orbit(orbit,E,K,model):
     """make an orbit"""
-      orbit.ee = E
-      #
-      # this should work, the boundaries are in radius...
-      orbit.r_circ = brentq(Ecirc,np.min(model.rcurve),np.max(model.rcurve),args=(orbit.ee,model))
-      orbit.kappa = K
-      orbit.jj = find_j(orbit.r_circ,orbit.kappa,model)
-      orbit.r_apo = brentq(denom,orbit.r_circ,np.max(model.rcurve),args=(orbit.ee,orbit.jj,model))
-      orbit.r_peri = brentq(denom,np.min(model.rcurve),orbit.r_circ,args=(orbit.ee,orbit.jj,model))
-      return orbit
+    orbit.ee = E
+    #
+    # this should work, the boundaries are in radius...
+    orbit.r_circ = brentq(Ecirc,np.min(model.rcurve),np.max(model.rcurve),args=(orbit.ee,model))
+    orbit.kappa = K
+    orbit.jj = find_j(orbit.r_circ,orbit.kappa,model)
+    orbit.r_apo = brentq(denom,orbit.r_circ,np.max(model.rcurve),args=(orbit.ee,orbit.jj,model))
+    orbit.r_peri = brentq(denom,np.min(model.rcurve),orbit.r_circ,args=(orbit.ee,orbit.jj,model))
+    return orbit
 
 
 

exptool/analysis/rotationcurves.py

@@ -2,6 +2,9 @@
 # tools to fit rotation curves. NEEDS MUCH CONSTRUCTION
 # from NewHorizon bars project
 
+import numpy as np
+astronomicalG = 0.0000043009125 # gravitational constant, (km/s)^2 * kpc / Msun
+
 
 def kuzmin_rotation(R,c,M,G=astronomicalG):
     """see BT08

exptool/observables/__init__.py

@@ -1,4 +1,14 @@
 
 
-# this is EXPtool
+__all__ = ['bulge',
+           'fourier',
+           'gaiatransform',
+           'movie',
+           'rotationcurves',
+           'scaleheighcheck',
+           'transform',
+           'velocity',
+           'visualize']
+
+
 

tests/importtest.py

@@ -0,0 +1,3 @@
+
+
+from exptool.analysis import *