diff --git a/.github/workflows/continuousintegration.yaml b/.github/workflows/continuousintegration.yaml index 8b183cb..dfef39c 100644 --- a/.github/workflows/continuousintegration.yaml +++ b/.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 diff --git a/exptool/analysis/__init__.py b/exptool/analysis/__init__.py index 18802fb..09346af 100644 --- a/exptool/analysis/__init__.py +++ b/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'] diff --git a/exptool/analysis/centering.py b/exptool/analysis/centering.py index 1cdb3d4..2674d8b 100644 --- a/exptool/analysis/centering.py +++ b/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)): """ diff --git a/exptool/analysis/resonancemodel.py b/exptool/analysis/resonancemodel.py index a66b94e..5456589 100644 --- a/exptool/analysis/resonancemodel.py +++ b/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 diff --git a/exptool/analysis/rotationcurves.py b/exptool/analysis/rotationcurves.py index 3b6a1f1..9767d83 100644 --- a/exptool/analysis/rotationcurves.py +++ b/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 diff --git a/exptool/observables/__init__.py b/exptool/observables/__init__.py index 2c28ec4..39eb753 100644 --- a/exptool/observables/__init__.py +++ b/exptool/observables/__init__.py @@ -1,4 +1,14 @@ -# this is EXPtool +__all__ = ['bulge', + 'fourier', + 'gaiatransform', + 'movie', + 'rotationcurves', + 'scaleheighcheck', + 'transform', + 'velocity', + 'visualize'] + + diff --git a/tests/importtest.py b/tests/importtest.py new file mode 100644 index 0000000..f8fa6b5 --- /dev/null +++ b/tests/importtest.py @@ -0,0 +1,3 @@ + + +from exptool.analysis import * \ No newline at end of file