"""
Simple units functionality for the openPMD beamphysics records.
For more advanced units, use a package like Pint:
https://pint.readthedocs.io/
"""
# import scipy.constants
m_e = 510998.94999999995
m_p = 938272088.16
c_light = 299792458
e_charge = 1.602176634e-19
mu_0 = 1.25663706212e-06
import numpy as np
[docs]
class pmd_unit:
"""
Params
------
unitSymbol: Native units name
unitSI: Conversion factor to the the correspontign SI unit
unitDimension: SI Base Exponents
Base unit dimensions are defined as:
Base dimension | exponents. | SI unit
---------------- ----------------- -------
length : (1,0,0,0,0,0,0) m
mass : (0,1,0,0,0,0,0) kg
time : (0,0,1,0,0,0,0) s
current : (0,0,0,1,0,0,0) A
temperture : (0,0,0,0,1,0,0) K
mol : (0,0,0,0,0,1,0) mol
luminous : (0,0,0,0,0,0,1) cd
Example:
pmd_unit('eV', 1.602176634e-19, (2, 1, -2, 0, 0, 0, 0))
defines that an eV is 1.602176634e-19 of base units m^2 kg/s^2, which is a Joule (J)
If unitSI=0 (default), init with a known symbol:
pmd_unit('T')
returns:
pmd_unit('T', 1, (0, 1, -2, -1, 0, 0, 0))
Simple equalities are provided:
u1 == u2
Returns True if the params are all the same.
"""
def __init__(self, unitSymbol="", unitSI=0, unitDimension=(0, 0, 0, 0, 0, 0, 0)):
# Allow to return an internally known unit
if unitSI == 0:
if unitSymbol in known_unit:
# Copy internals
u = known_unit[unitSymbol]
unitSI = u.unitSI
unitDimension = u.unitDimension
else:
raise ValueError(f"unknown unitSymbol: {unitSymbol}")
self._unitSymbol = unitSymbol
self._unitSI = unitSI
if isinstance(unitDimension, str):
self._unitDimension = DIMENSION[unitDimension]
else:
self._unitDimension = unitDimension
@property
def unitSymbol(self):
return self._unitSymbol
@property
def unitSI(self):
return self._unitSI
@property
def unitDimension(self):
return self._unitDimension
def __mul__(self, other):
return multiply_units(self, other)
def __truediv__(self, other):
return divide_units(self, other)
def __eq__(self, other):
if isinstance(other, self.__class__):
return self.__dict__ == other.__dict__
else:
return False
def __ne__(self, other):
return not self.__eq__(other)
def __str__(self):
return self.unitSymbol
def __repr__(self):
return f"pmd_unit('{self.unitSymbol}', {self.unitSI}, {self.unitDimension})"
[docs]
def is_identity(u):
"""Checks if the unit is equivalent to 1"""
return u.unitSI == 1 and u.unitDimension == (0, 0, 0, 0, 0, 0, 0)
[docs]
def multiply_units(u1, u2):
"""
Multiplies two pmd_unit symbols
"""
if is_identity(u1):
return u2
if is_identity(u2):
return u1
s1 = u1.unitSymbol
s2 = u2.unitSymbol
if s1 == s2:
symbol = f"{s1}^2"
else:
symbol = s1 + "*" + s2
d1 = u1.unitDimension
d2 = u2.unitDimension
dim = tuple(sum(x) for x in zip(d1, d2))
unitSI = u1.unitSI * u2.unitSI
return pmd_unit(unitSymbol=symbol, unitSI=unitSI, unitDimension=dim)
[docs]
def divide_units(u1, u2):
"""
Divides two pmd_unit symbols : u1/u2
"""
if is_identity(u2):
return u1
s1 = u1.unitSymbol
s2 = u2.unitSymbol
if s1 == s2:
symbol = "1"
else:
symbol = s1 + "/" + s2
d1 = u1.unitDimension
d2 = u2.unitDimension
dim = tuple(a - b for a, b in zip(d1, d2))
unitSI = u1.unitSI / u2.unitSI
return pmd_unit(unitSymbol=symbol, unitSI=unitSI, unitDimension=dim)
[docs]
def sqrt_unit(u):
"""
Returns the sqrt of a unit
"""
u.unitDimension
symbol = u.unitSymbol
if symbol not in ["", "1"]:
symbol = f"sqrt({symbol})"
unitSI = np.sqrt(u.unitSI)
dim = tuple(x / 2 for x in u.unitDimension)
return pmd_unit(unitSymbol=symbol, unitSI=unitSI, unitDimension=dim)
DIMENSION = {
"1": (0, 0, 0, 0, 0, 0, 0),
# Base units
"length": (1, 0, 0, 0, 0, 0, 0),
"mass": (0, 1, 0, 0, 0, 0, 0),
"time": (0, 0, 1, 0, 0, 0, 0),
"current": (0, 0, 0, 1, 0, 0, 0),
"temperture": (0, 0, 0, 0, 1, 0, 0),
"mol": (0, 0, 0, 0, 0, 1, 0),
"luminous": (0, 0, 0, 0, 0, 0, 1),
#
"charge": (0, 0, 1, 1, 0, 0, 0),
"electric_field": (1, 1, -3, -1, 0, 0, 0),
"electric_potential": (1, 2, -3, -1, 0, 0, 0),
"magnetic_field": (0, 1, -2, -1, 0, 0, 0),
"velocity": (1, 0, -1, 0, 0, 0, 0),
"energy": (2, 1, -2, 0, 0, 0, 0),
"momentum": (1, 1, -1, 0, 0, 0, 0),
}
# Inverse
DIMENSION_NAME = {v: k for k, v in DIMENSION.items()}
[docs]
def dimension(name):
if name in DIMENSION:
return DIMENSION[name]
else:
return None
[docs]
def dimension_name(dim_array):
return DIMENSION_NAME[tuple(dim_array)]
SI_symbol = {
"1": "1",
"length": "m",
"mass": "kg",
"time": "s",
"current": "A",
"temperture": "K",
"mol": "mol",
"luminous": "cd",
"charge": "C",
"electric_field": "V/m",
"electric_potential": "V",
"velocity": "m/s",
"energy": "J",
"momentum": "kg*m/s",
"magnetic_field": "T",
}
# Inverse
SI_name = {v: k for k, v in SI_symbol.items()}
known_unit = {
"1": pmd_unit("", 1, "1"),
"rad": pmd_unit("", 1, "1"),
"m": pmd_unit("m", 1, "length"),
"kg": pmd_unit("kg", 1, "mass"),
"g": pmd_unit("g", 0.001, "mass"),
"s": pmd_unit("s", 1, "time"),
"A": pmd_unit("A", 1, "current"),
"K": pmd_unit("K", 1, "temperture"),
"mol": pmd_unit("mol", 1, "mol"),
"cd": pmd_unit("cd", 1, "luminous"),
"C": pmd_unit("C", 1, "charge"),
"charge_num": pmd_unit("charge #", 1, "charge"),
"V/m": pmd_unit("V/m", 1, "electric_field"),
"V": pmd_unit("V", 1, "electric_potential"),
"c_light": pmd_unit("vel/c", c_light, "velocity"),
"c": pmd_unit("vel/c", c_light, "velocity"),
"m/s": pmd_unit("m/s", 1, "velocity"),
"eV": pmd_unit("eV", e_charge, "energy"),
"J": pmd_unit("J", 1, "energy"),
"eV/c": pmd_unit("eV/c", e_charge / c_light, "momentum"),
"T": pmd_unit("T", 1, "magnetic_field"),
}
[docs]
def unit(symbol):
"""
Returns a pmd_unit from a known symbol.
* is allowed between two known symbols:
"""
if symbol in known_unit:
return known_unit[symbol]
if "*" in symbol:
subunits = [known_unit[s] for s in symbol.split("*")]
# Require these to be in known units
assert len(subunits) == 2, "TODO: more complicated units"
return multiply_units(subunits[0], subunits[1])
raise ValueError(f"Unknown unit symbol: {symbol}")
# Dicts for prefixes
PREFIX_FACTOR = {
"yocto-": 1e-24,
"zepto-": 1e-21,
"atto-": 1e-18,
"femto-": 1e-15,
"pico-": 1e-12,
"nano-": 1e-9,
"micro-": 1e-6,
"milli-": 1e-3,
"centi-": 1e-2,
"deci-": 1e-1,
"deca-": 1e1,
"hecto-": 1e2,
"kilo-": 1e3,
"mega-": 1e6,
"giga-": 1e9,
"tera-": 1e12,
"peta-": 1e15,
"exa-": 1e18,
"zetta-": 1e21,
"yotta-": 1e24,
}
# Inverse
PREFIX = dict((v, k) for k, v in PREFIX_FACTOR.items())
SHORT_PREFIX_FACTOR = {
"y": 1e-24,
"z": 1e-21,
"a": 1e-18,
"f": 1e-15,
"p": 1e-12,
"n": 1e-9,
"ยต": 1e-6,
"m": 1e-3,
"c": 1e-2,
"d": 1e-1,
"": 1,
"da": 1e1,
"h": 1e2,
"k": 1e3,
"M": 1e6,
"G": 1e9,
"T": 1e12,
"P": 1e15,
"E": 1e18,
"Z": 1e21,
"Y": 1e24,
}
# Inverse
SHORT_PREFIX = dict((v, k) for k, v in SHORT_PREFIX_FACTOR.items())
# Nice scaling
[docs]
def nice_scale_prefix(scale):
"""
Returns a nice factor and a SI prefix string
Example:
scale = 2e-10
f, u = nice_scale_prefix(scale)
"""
if scale == 0:
return 1, ""
p10 = np.log10(abs(scale))
if p10 < -2 or p10 > 2:
f = 10 ** (p10 // 3 * 3)
else:
f = 1
return f, SHORT_PREFIX[f]
[docs]
def nice_array(a):
"""
Returns a scaled array, the scaling, and a unit prefix
Example:
nice_array( np.array([2e-10, 3e-10]) )
Returns:
(array([200., 300.]), 1e-12, 'p')
"""
if np.isscalar(a):
x = a
elif len(a) == 1:
x = a[0]
else:
a = np.array(a)
x = a.ptp()
fac, prefix = nice_scale_prefix(x)
return a / fac, fac, prefix
# -------------------------
# Units for ParticleGroup
PARTICLEGROUP_UNITS = {}
for k in ["status"]:
PARTICLEGROUP_UNITS[k] = unit("1")
for k in ["t"]:
PARTICLEGROUP_UNITS[k] = unit("s")
for k in [
"energy",
"kinetic_energy",
"mass",
"higher_order_energy_spread",
"higher_order_energy",
]:
PARTICLEGROUP_UNITS[k] = unit("eV")
for k in ["px", "py", "pz", "p", "pr"]:
PARTICLEGROUP_UNITS[k] = unit("eV/c")
for k in ["x", "y", "z", "r", "Jx", "Jy"]:
PARTICLEGROUP_UNITS[k] = unit("m")
for k in ["beta", "beta_x", "beta_y", "beta_z", "gamma", "theta", "ptheta"]:
PARTICLEGROUP_UNITS[k] = unit("1")
for k in ["charge", "species_charge", "weight"]:
PARTICLEGROUP_UNITS[k] = unit("C")
for k in ["average_current"]:
PARTICLEGROUP_UNITS[k] = unit("A")
for k in ["norm_emit_x", "norm_emit_y"]:
PARTICLEGROUP_UNITS[k] = unit("m")
for k in ["norm_emit_4d"]:
PARTICLEGROUP_UNITS[k] = multiply_units(unit("m"), unit("m"))
for k in ["xp", "yp"]:
PARTICLEGROUP_UNITS[k] = unit("rad")
for k in ["x_bar", "px_bar", "y_bar", "py_bar"]:
PARTICLEGROUP_UNITS[k] = sqrt_unit(unit("m"))
[docs]
def pg_units(key):
"""
Returns a str representing the units of any attribute
"""
for prefix in ["sigma_", "mean_", "min_", "max_", "ptp_", "delta_"]:
if key.startswith(prefix):
nkey = key[len(prefix) :]
return PARTICLEGROUP_UNITS[nkey]
if key.startswith("cov_"):
subkeys = key.strip("cov_").split("__")
unit0 = PARTICLEGROUP_UNITS[subkeys[0]]
unit1 = PARTICLEGROUP_UNITS[subkeys[1]]
return multiply_units(unit0, unit1)
# Fields
if key.startswith("electricField"):
return unit("V/m")
if key.startswith("magneticField"):
return unit("T")
return PARTICLEGROUP_UNITS[key]
# -------------------------
# h5 tools
[docs]
def write_unit_h5(h5, u):
"""
Writes an pmd_unit to an h5 handle
"""
h5.attrs["unitSI"] = u.unitSI
h5.attrs["unitDimension"] = u.unitDimension
h5.attrs["unitSymbol"] = u.unitSymbol
[docs]
def read_unit_h5(h5):
"""
Reads unit data from an h5 handle and returns a pmd_unit object
"""
a = h5.attrs
unitSI = a["unitSI"]
unitDimension = tuple(a["unitDimension"])
if "unitSymbol" not in a:
unitSymbol = "unknown"
else:
unitSymbol = a["unitSymbol"]
return pmd_unit(unitSymbol=unitSymbol, unitSI=unitSI, unitDimension=unitDimension)
[docs]
def read_dataset_and_unit_h5(h5, expected_unit=None, convert=True):
"""
Reads a dataset that has openPMD unit attributes.
expected_unit can be a pmd_unit object, or a known unit str. Examples: 'kg', 'J', 'eV'
If expected_unit is given, will check that the units are compatible.
If convert, the data will be returned with the expected_units.
Returns a tuple:
np.array, pmd_unit
"""
# Read the unit that is there.
u = read_unit_h5(h5)
# Simple case
if not expected_unit:
return np.array(h5), u
if isinstance(expected_unit, str):
# Try to get unit
expected_unit = unit(expected_unit)
# Check dimensions
du = divide_units(u, expected_unit)
assert du.unitDimension == (0, 0, 0, 0, 0, 0, 0), "incompatible units"
if convert:
fac = du.unitSI
return fac * np.array(h5), expected_unit
else:
return np.array(h5), u
[docs]
def write_dataset_and_unit_h5(h5, name, data, unit=None):
"""
Writes data and pmd_unit to h5[name]
See: read_dataset_and_unit_h5
"""
h5[name] = data
if unit:
write_unit_h5(h5[name], unit)