From c69f5e462fb5a40f30535c744b18c958df016638 Mon Sep 17 00:00:00 2001 From: David Minton Date: Mon, 26 Feb 2024 10:50:03 -0500 Subject: [PATCH] Refactored argument name that rotates the coordinate frame to align with the central body rotation and now explicitly pass it to the internal methods explicitly to prevent an error in the JPL Horizons call --- swiftest/simulation_class.py | 35 ++++++++++++++++++++++------------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/swiftest/simulation_class.py b/swiftest/simulation_class.py index 688dc07fa..61f54d64a 100644 --- a/swiftest/simulation_class.py +++ b/swiftest/simulation_class.py @@ -2150,6 +2150,7 @@ def add_solar_system_body(self, ephemeris_id: int | List[int] | None = None, date: str | None = None, source: str = "HORIZONS", + align_to_central_body_rotation: bool = False, **kwargs: Any ) -> None: """ @@ -2178,7 +2179,10 @@ def add_solar_system_body(self, set by `set_ephemeris_date`. source : str, default "Horizons" The source of the ephemerides. - Currently only the JPL Horizons ephemeris is implemented, so this is ignored. + Currently only the JPL Horizons ephemeris is implemented, so this is ignored. + align_to_central_body_rotation : bool, default False + If True, the cartesian coordinates will be aligned to the rotation pole of the central body. This is only valid for when + rotation is enabled. **kwargs : Any Additional keyword arguments to pass to the query method (i.e. astroquery.Horizons) @@ -2265,7 +2269,7 @@ def add_solar_system_body(self, dsnew = init_cond.vec2xr(self.param,**kwargs) - dsnew = self._combine_and_fix_dsnew(dsnew,**kwargs) + dsnew = self._combine_and_fix_dsnew(dsnew,align_to_central_body_rotation, **kwargs) if dsnew['id'].max(dim='name') > 0 and dsnew['name'].size > 0: self.save(verbose=False) @@ -2422,6 +2426,7 @@ def add_body(self, J4: float | List[float] | npt.NDArray[np.float_] | None=None, c_lm: List[float] | List[npt.NDArray[np.float_]] | npt.NDArray[np.float_] | None = None, rotphase: float | List[float] | npt.NDArray[np.float_] | None=None, + align_to_central_body_rotation: bool = False, **kwargs: Any ) -> None: """ @@ -2467,7 +2472,9 @@ def add_body(self, Principal axes moments of inertia vectors if these are massive bodies with rotation enabled. rotphase : float, optional rotation phase angle of the central body in degrees - + align_to_central_body_rotation : bool, default False + If True, the cartesian coordinates will be aligned to the rotation pole of the central body. This is only valid for when + rotation is enabled. Returns ------- None @@ -2602,13 +2609,14 @@ def input_to_clm_array(val, n): dsnew = init_cond.vec2xr(self.param, name=name, a=a, e=e, inc=inc, capom=capom, omega=omega, capm=capm, id=id, Gmass=Gmass, radius=radius, rhill=rhill, Ip=Ip, rh=rh, vh=vh,rot=rot, j2rp2=J2, j4rp4=J4, c_lm=c_lm, rotphase=rotphase, time=time) - dsnew = self._combine_and_fix_dsnew(dsnew,**kwargs) + dsnew = self._combine_and_fix_dsnew(dsnew,align_to_central_body_rotation,**kwargs) self.save(verbose=False) return def _combine_and_fix_dsnew(self, dsnew: xr.Dataset, + align_to_central_body_rotation: bool = False, **kwargs: Any ) -> xr.Dataset: """ @@ -2618,6 +2626,9 @@ def _combine_and_fix_dsnew(self, ---------- dsnew : xarray Dataset Dataset with new bodies + align_to_central_body_rotation : bool, default False + If True, the cartesian coordinates will be aligned to the rotation pole of the central body. This is only valid for when + rotation is enabled. Returns ------- @@ -2644,7 +2655,7 @@ def _combine_and_fix_dsnew(self, dsnew = io.fix_types(dsnew, ftype=np.float32) self.data = io.fix_types(self.data, ftype=np.float32) - self.set_central_body(**kwargs) + self.set_central_body(align_to_central_body_rotation) def get_nvals(ds): if "name" in ds.dims: count_dim = "name" @@ -3190,14 +3201,14 @@ def clean(self): return def set_central_body(self, - align_to_rotation_pole: bool = False, + align_to_central_body_rotation: bool = False, **kwargs: Any): """ - Sets the central body to be the most massive body in the dataset. Cartesian position and velocity Cartesian coordinates are rotated If align_to_rotation_pole is True, the rotation pole is set to the z-axis. + Sets the central body to be the most massive body in the dataset. Cartesian position and velocity Cartesian coordinates are rotated If align_to_central_body_rotation is True, the rotation pole is set to the z-axis. Parameters ---------- - align_to_rotation_pole : bool, default False + align_to_central_body_rotation : bool, default False If True, the rotation pole is set to the z-axis. Returns @@ -3206,7 +3217,6 @@ def set_central_body(self, """ - if "Gmass" not in self.data: warnings.warn("No bodies with Gmass values found in dataset. Cannot set central body.",stacklevel=2) return @@ -3240,11 +3250,10 @@ def set_central_body(self, pos_skip = ['space','Ip','rot'] for var in self.data.variables: - if var not in pos_skip: + if 'space' in self.data[var].dims and var not in pos_skip: self.data[var] -= cbda[var] - if align_to_rotation_pole and 'rot' in cbda: - self.data = tool.rotate_to_vector(self.data,cbda.rot) + if align_to_central_body_rotation and 'rot' in cbda: + self.data = tool.rotate_to_vector(self.data,cbda.rot.isel(time=0).values[()]) - return \ No newline at end of file