From b04572046d659f400d27020d321ada5e51d965b6 Mon Sep 17 00:00:00 2001 From: nstarman Date: Fri, 10 Nov 2023 12:14:12 -0500 Subject: [PATCH] cleanup Signed-off-by: nstarman --- noxfile.py | 5 +---- src/galdynamix/dynamics/_orbit.py | 2 +- src/galdynamix/dynamics/mockstream/_core.py | 2 +- src/galdynamix/dynamics/mockstream/_df/fardal.py | 10 ++++++---- .../dynamics/mockstream/_mockstream_generator.py | 7 +------ src/galdynamix/integrate/_builtin.py | 6 ++---- src/galdynamix/potential/_potential/base.py | 3 +-- 7 files changed, 13 insertions(+), 22 deletions(-) diff --git a/noxfile.py b/noxfile.py index aadab843..d8121b2a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -46,10 +46,7 @@ def docs(session: nox.Session) -> None: parser = argparse.ArgumentParser() parser.add_argument("--serve", action="store_true", help="Serve after building") parser.add_argument( - "-b", - dest="builder", - default="html", - help="Build target (default: html)", + "-b", dest="builder", default="html", help="Build target (default: html)" ) args, posargs = parser.parse_known_args(session.posargs) diff --git a/src/galdynamix/dynamics/_orbit.py b/src/galdynamix/dynamics/_orbit.py index a6c41fd5..9d63fe4d 100644 --- a/src/galdynamix/dynamics/_orbit.py +++ b/src/galdynamix/dynamics/_orbit.py @@ -37,7 +37,7 @@ class Orbit(eqx.Module): # type: ignore[misc] @property @partial_jit() def qp(self) -> jt.Array: - """Return as a single Array[(N, Q + P + T),].""" + """Return as a single Array[(N, Q + P),].""" # Determine output shape qd = self.q.shape[1] # dimensionality of q shape = (self.q.shape[0], qd + self.p.shape[1]) diff --git a/src/galdynamix/dynamics/mockstream/_core.py b/src/galdynamix/dynamics/mockstream/_core.py index 9d728261..b4f3f874 100644 --- a/src/galdynamix/dynamics/mockstream/_core.py +++ b/src/galdynamix/dynamics/mockstream/_core.py @@ -35,7 +35,7 @@ class MockStream(eqx.Module): # type: ignore[misc] @property @partial_jit() def qp(self) -> jt.Array: - """Return as a single Array[(N, Q + P + T),].""" + """Return as a single Array[(N, Q + P),].""" # Determine output shape qd = self.q.shape[1] # dimensionality of q shape = (self.q.shape[0], qd + self.p.shape[1]) diff --git a/src/galdynamix/dynamics/mockstream/_df/fardal.py b/src/galdynamix/dynamics/mockstream/_df/fardal.py index b9d9699b..e8b1a4c5 100644 --- a/src/galdynamix/dynamics/mockstream/_df/fardal.py +++ b/src/galdynamix/dynamics/mockstream/_df/fardal.py @@ -81,8 +81,9 @@ def _sample( kvz_samp = kvz_bar + jax.random.normal(keyd, shape=(1,)) * sigma_kvz # Trailing arm - x_trail = x + kr_samp * r_hat * (r_tidal) # nudge out - x_trail = x_trail + z_hat * kz_samp * (r_tidal / 1.0) + x_trail = ( + x + (kr_samp * r_hat * (r_tidal)) + (z_hat * kz_samp * (r_tidal / 1.0)) + ) v_trail = ( v + (0.0 + kvphi_samp * v_circ * (1.0)) * phi_hat @@ -90,8 +91,9 @@ def _sample( ) # Leading arm - x_lead = x + kr_samp * r_hat * (-r_tidal) # nudge in - x_lead = x_lead + z_hat * kz_samp * (-r_tidal / 1.0) + x_lead = ( + x + (kr_samp * r_hat * (-r_tidal)) + (z_hat * kz_samp * (-r_tidal / 1.0)) + ) v_lead = ( v + (0.0 + kvphi_samp * v_circ * (-1.0)) * phi_hat diff --git a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py b/src/galdynamix/dynamics/mockstream/_mockstream_generator.py index 06b5f249..73e3a992 100644 --- a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py +++ b/src/galdynamix/dynamics/mockstream/_mockstream_generator.py @@ -30,12 +30,7 @@ class MockStreamGenerator(eqx.Module): # type: ignore[misc] @partial_jit(static_argnames=("seed_num",)) def _run_scan( - self, - ts: jt.Array, - prog_w0: jt.Array, - prog_mass: jt.Array, - *, - seed_num: int, + self, ts: jt.Array, prog_w0: jt.Array, prog_mass: jt.Array, *, seed_num: int ) -> tuple[tuple[jt.Array, jt.Array], Orbit]: """Generate stellar stream by scanning over the release model/integration. diff --git a/src/galdynamix/integrate/_builtin.py b/src/galdynamix/integrate/_builtin.py index f662ce15..98d6dd75 100644 --- a/src/galdynamix/integrate/_builtin.py +++ b/src/galdynamix/integrate/_builtin.py @@ -28,8 +28,7 @@ class DiffraxIntegrator(AbstractIntegrator): Solver: AbstractSolver = eqx.field(default=Dopri5, static=True) SaveAt: DiffraxSaveAt = eqx.field(default=DiffraxSaveAt, static=True) stepsize_controller: AbstractStepSizeController = eqx.field( - default=PIDController(rtol=1e-7, atol=1e-7), - static=True, + default=PIDController(rtol=1e-7, atol=1e-7), static=True ) diffeq_kw: tuple[tuple[str, Any], ...] = eqx.field( default_factory=lambda: ( @@ -39,8 +38,7 @@ class DiffraxIntegrator(AbstractIntegrator): static=True, ) solver_kw: tuple[tuple[str, Any], ...] = eqx.field( - default_factory=lambda: (("scan_kind", "bounded"),), - static=True, + default_factory=lambda: (("scan_kind", "bounded"),), static=True ) def run( diff --git a/src/galdynamix/potential/_potential/base.py b/src/galdynamix/potential/_potential/base.py index 2a5a81c9..33c31832 100644 --- a/src/galdynamix/potential/_potential/base.py +++ b/src/galdynamix/potential/_potential/base.py @@ -51,8 +51,7 @@ def _init_units(self) -> None: value = getattr(self, f.name) if isinstance(value, u.Quantity): value = value.to_value( - self.units[param.physical_type], - equivalencies=param.equivalencies, + self.units[param.physical_type], equivalencies=param.equivalencies ) object.__setattr__(self, f.name, value)