def test_write(): """ Tests whether swiftsimio can handle a new particle type. If the test doesn't crash this is a success. """ # Use default units, i.e. cm, grams, seconds, Ampere, Kelvin unit_system = unyt.UnitSystem(name="default", length_unit=unyt.cm, mass_unit=unyt.g, time_unit=unyt.s) # Specify a new type in the metadata - currently done by editing the dictionaries directly. # TODO: Remove this terrible way of setting up different particle types. swp.particle_name_underscores[6] = "extratype" swp.particle_name_class[6] = "Extratype" swp.particle_name_text[6] = "Extratype" swmw.extratype = {"smoothing_length": "SmoothingLength", **swmw.shared} boxsize = 10 * unyt.cm x = Writer(unit_system, boxsize, unit_fields_generate_units=generate_units) x.extratype.coordinates = np.zeros((10, 3)) * unyt.cm for i in range(0, 10): x.extratype.coordinates[i][0] = float(i) x.extratype.velocities = np.zeros((10, 3)) * unyt.cm / unyt.s x.extratype.masses = np.ones(10, dtype=float) * unyt.g x.extratype.smoothing_length = np.ones(10, dtype=float) * (5.0 * unyt.cm) x.write("extra_test.hdf5") # Clean up these global variables we screwed around with... swp.particle_name_underscores.pop(6) swp.particle_name_class.pop(6) swp.particle_name_text.pop(6)
# generate the coordinates dist = np.random.rand(num_part) * (max_r - min_r) + min_r angle = np.random.rand(num_part) * 2 * np.pi # more easy to do it in 2D, therefore coords[:, 2] == 0 coords = np.zeros((num_part, 3)) coords[:, 0] = dist * np.sin(angle) coords[:, 1] = dist * np.cos(angle) coords += boxsize * 0.5 # generate the masses m = np.ones(num_part) * masses # generate the velocities sign = np.random.rand(num_part) sign[sign < 0.5] = -1 sign[sign >= 0.5] = 1 v = np.zeros((num_part, 3)) v[:, 0] = sign * np.sqrt(G * M / (dist * (1 + np.tan(angle)**2))) v[:, 1] = - np.tan(angle) * v[:, 0] # Write the snapshot units = unyt.UnitSystem("Planets", unyt.AU, unyt.mearth, unyt.yr) snapshot = Writer(units, boxsize * unyt.AU) snapshot.dark_matter.coordinates = coords * unyt.AU snapshot.dark_matter.velocities = v * unyt.AU / unyt.yr snapshot.dark_matter.masses = m * unyt.mearth snapshot.write(filename)
""" Contains unit systems that may be useful to astronomers. In particular, it contains the cosmo_units which can be considered Gadget-oid default units, with + Unit length = Mpc + Unit velocity = km/s + Unit mass = 10^10 Msun + Unit temperature = K """ import unyt cosmo_units = unyt.UnitSystem( "cosmological", unyt.Mpc, 1e10 * unyt.msun, (1.0 * unyt.s * unyt.Mpc / unyt.km).to(unyt.Gyr), )
"""Define stress/strain tensors.""" from numpy import array, empty import unyt as u from unyt.array import unyt_array # set unit regsitry us_unit_system = u.UnitSystem(name='US', length_unit='inch', mass_unit='lb', time_unit='s', temperature_unit='degF', angle_unit='deg') # add strain unit ureg = u.UnitRegistry() ureg.add('strain', base_value=1.0, dimensions=u.dimensions.dimensionless, tex_repr=r"\rm{\varepsilon}", prefixable=True) u.unit_registry.default_unit_registry = ureg class Tensor(unyt_array): """Base 3x3 symmetric tensor. Parameters ---------- matrix : numpy.array_like an array-like object with six values
""" Contains unit systems that may be useful to astronomers. In particular, it contains the cosmo_units which can be considered Gadget-oid default units, with + Unit length = Mpc + Unit velocity = km/s + Unit mass = 10^10 Msun + Unit temperature = K """ import unyt cosmo_units = unyt.UnitSystem( "cosmological", unyt.Mpc, 1e10 * unyt.msun, unyt.s * unyt.Mpc / unyt.km )
"""Utility functions for handling units.""" import unyt as u planckton_units = u.UnitSystem("planckton", mass_unit="amu", length_unit="nm", time_unit="s") planckton_units["energy"] = "kJ" constants = { "avogadro": 6.022140857e23 / u.mol, "boltzmann": 1.38064852e-23 * u.J / u.K, } def quantity_to_tuple(quantity): """Break a unyt.quantity into a tuple. Convert a unyt.quantity into a tuple containing its value and units in string format. Useful for serialization. IMPORTANT: This function expects one quantity, not an array. Parameters ---------- quantity: unyt.unyt_quantity Returns ------- (number, string)
""" from pygadgetreader import * from swiftsimio import Writer import numpy as np import unyt filename = "/cosma7/data/dp004/jlvc76/nIFTy/IC_CLUSTER_00019" length = unyt.kpc mass = 1e10 * unyt.msun time = (1.0 * unyt.s * unyt.kpc / unyt.km).to("s") velocity = length / time energy_per_unit_mass = (length / time)**2 nifty_units = unyt.UnitSystem("nifty", 1e3 * length, mass, time) writer = Writer( unit_system=nifty_units, box_size=readheader(filename, "boxsize") * length, dimension=3, compress=True, extra_header={ "Redshift": readheader(filename, "redshift"), "Omega0": readheader(filename, "O0"), "OmegaLambda": readheader(filename, "Ol"), "HubbleParam": readheader(filename, "h"), }, ) writer.gas.coordinates = unyt.unyt_array(readsnap(filename, "pos", 0), length)
def create_unit_system(length: UnitLike, mass: UnitLike, time: UnitLike, temperature: UnitLike = None, angle: UnitLike = None, current_mks: UnitLike = None, luminous_intensity: UnitLike = None, logarithmic: UnitLike = None, name: str = None, registry: t.Optional[unyt.UnitRegistry] = None, consistent: bool = False, strict_dims: bool = True, **convenience_units: UnitLike) -> unyt.UnitSystem: """ Create a new unit system. Parameters ---------- length : UnitLike The base length unit. mass : UnitLike The base mass unit. time : UnitLike The base time unit. temperature : UnitLike, optional The base temperature unit. (default: 'K') angle : UnitLike, optional The base angle unit. (default: 'rad') current_mks : UnitLike, optional The base current unit. (default: 'A') luminous_intensity : UnitLike, optional The base luminous intensity unit. (default: 'Cd') logarithmic : UnitLike, optional The base logarithmic unit. (default: 'Np') name : str, optional Name for the unit system. If not provided, a name is generated from the provided base units. (default: None) registry : UnitRegistry, optional The unit registry for the system. If None, the default unit registry is used. consistent : bool, optional If True, enforce consistency between convenience units and base units such that, in base units, convenience units have a magnitude of 1.0. For example, force='N' is consistent with MKS (N = 1.0 kg*m/s**2), but force='kN' is not (kN = 1e3 kg*m/s**2). (default: False) strict_dims : bool, optional If True, strictly enforce convenience units having the dimensions they map to. (default: True) **convenience_units : str, optional Mapping of dimension names to convenience units, e.g., ``force='kN'``. Raises ------ IllDefinedUnitSystem If `strict_dims` is True and convenience units do not strictly map UnitSystemExistsError If a unit system with name `name` already exists UnitSystemConsistencyError If `consistent` is True and unit system is not consistent Example ------- Unit system that uses millimeters and gigagrams, with a convenience kilonewton unit: >>> system = create_unit_system('mm', 'Gg', 's', force='kN') >>> system mm_Gg_s Unit System Base Units: length: mm mass: Gg time: s temperature: K angle: rad current_mks: A luminous_intensity: cd logarithmic: Np Other Units: force: kN Unit system with different base temperature unit, but still auto-generated name: >>> system = create_unit_system('m', 'kg', 's', temperature='degC') >>> system m_kg_s_degC Unit System Base Units: length: m mass: kg time: s temperature: degC angle: rad current_mks: A luminous_intensity: cd logarithmic: Np Other Units: """ base_units = dict(length_unit=length, mass_unit=mass, time_unit=time) # Handle other default base units; don't include them in auto-generated name # if not provided. They already have defaults set by the UnitSystem # constructor. if temperature is not None: base_units['temperature_unit'] = temperature if angle is not None: base_units['angle_unit'] = angle if current_mks is not None: base_units['current_mks_unit'] = current_mks if luminous_intensity is not None: base_units['luminous_intensity_unit'] = luminous_intensity if logarithmic is not None: base_units['logarithmic_unit'] = logarithmic # Generate name if not provided. if name is None: name = '_'.join(map(str, base_units.values())) # Check against existing unit systems. if name in unyt.unit_systems.unit_system_registry: raise UnitSystemExistsError(name) if registry is None: registry = unyt.unit_registry.default_unit_registry # Create new system with base units. system = unyt.UnitSystem(str(name), **base_units, registry=registry) try: # Apply convenience units. for dim, unit in convenience_units.items(): if strict_dims: dim_obj = getattr(unyt.dimensions, dim) unit_obj = unyt.Unit(unit) if unit_obj.dimensions != dim_obj: raise IllDefinedUnitSystem( f'{dim} [{dim_obj}] -> {unit_obj} ' f'[{unit_obj.dimensions}]') system[dim] = unit # Check consistency if asked to. if consistent: check = check_consistent_unit_system(system) if not check.is_consistent: raise UnitSystemConsistencyError( f'Inconsistent unit for dimension {check.bad_dim!r}: ' f'1.0 {check.bad_unit} = {check.bad_base}') except: # Remove bad system from registry del unyt.unit_systems.unit_system_registry[name] raise return system
""" Contains unit systems that may be useful to astronomers. In particular, it contains the cosmo_units which can be considered Gadget-oid default units, with + Unit length = Mpc + Unit velocity = km/s + Unit mass = 10^10 Msun + Unit temperature = K Also contains unit conversion factors, to simplify units wherever possible. """ import unyt try: # Need to do this first otherwise the `unyt` system freaks out about # us upgrading msun from a symbol cosmo_units = unyt.UnitSystem( "cosmological", unyt.Mpc, unyt.unyt_quantity(1e10, units=unyt.Solar_Mass), unyt.unyt_quantity(1.0, units=unyt.s * unyt.Mpc / unyt.km).to(unyt.Gyr), ) except RuntimeError: # We've already done that, oops. cosmo_units = unyt.unit_systems.cosmological pass