42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import astropy.units as u
|
|
import numpy as np
|
|
import logging
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
M_SCALE: int = None
|
|
R_SCALE: int = None
|
|
|
|
def seed_scales(r_scale: u.Quantity, m_scale: u.Quantity):
|
|
"""
|
|
Set the scales for the given simulation.
|
|
Parameters:
|
|
- r_scale: astropy.units.Quantity with units of length - the characteristic length scale of the simulation. Particle positions are expressed in units of this scale.
|
|
- m_scale: astropy.units.Quantity with units of mass - the characteristic mass scale of the simulation. Particle masses are expressed in units of this scale.
|
|
"""
|
|
global M_SCALE, R_SCALE
|
|
M_SCALE = m_scale
|
|
R_SCALE = r_scale
|
|
logger.info(f"Set scales: M_SCALE = {M_SCALE}, R_SCALE = {R_SCALE}")
|
|
|
|
|
|
def apply_units(columns: np.array, quantity: str):
|
|
if quantity == "mass":
|
|
return columns * M_SCALE
|
|
elif quantity == "position":
|
|
return columns * R_SCALE
|
|
elif quantity == "velocity":
|
|
return columns * R_SCALE / u.s
|
|
elif quantity == "time":
|
|
return columns * u.s
|
|
elif quantity == "acceleration":
|
|
return columns * R_SCALE / u.s**2
|
|
elif quantity == "force":
|
|
return columns * M_SCALE * R_SCALE / u.s**2
|
|
elif quantity == "volume":
|
|
return columns * R_SCALE**3
|
|
elif quantity == "density":
|
|
return columns * M_SCALE / R_SCALE**3
|
|
else:
|
|
raise ValueError(f"Unknown quantity: {quantity}")
|