Source code for qcodes.math_utils.field_vector

"""
A helper module containing a class to keep track of vectors in different
coordinate systems.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeVar

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

AllCoordsType = tuple[float, float, float, float, float, float, float]
NormOrder = None | float | Literal["fro"] | Literal["nuc"]
T = TypeVar("T", bound="FieldVector")


[docs] class FieldVector: """ A convenient class to keep track of vectors representing physical fields. The idea is that a vector instance stores a representation in Cartesian, spherical and cylindrical coordinates. All arguments are optional, however the user needs to provide one of the combinations of either (x, y, z) values or (rho, phi, z) values or (r, theta, phi) values at instantiation for a meaningful computation of the other representation, immediately. """ attributes: ClassVar[list[str]] = ["x", "y", "z", "r", "theta", "phi", "rho"] repr_format = "cartesian" def __init__( self, x: float | None = None, y: float | None = None, z: float | None = None, r: float | None = None, theta: float | None = None, phi: float | None = None, rho: float | None = None, ): """ Args: x: represents the norm of the projection of the vector along the x-axis y: represents the norm of the projection of the vector along the y-axis z: represents the norm of the projection of the vector along the z-axis r: represents the norm of the vector theta: represents the angle of the vector with respect to the positive z-axis rho: represents the norm of the projection of the vector on to the xy-plane phi: represents the angle of rho with respect to the positive x-axis """ self._x = float(x) if x is not None else None self._y = float(y) if y is not None else None self._z = float(z) if z is not None else None self._r = float(r) if r is not None else None self._theta = float(np.radians(theta)) if theta is not None else None self._phi = float(np.radians(phi)) if phi is not None else None self._rho = float(rho) if rho is not None else None self._compute_unknowns() def _set_attribute_value(self, attr_name: str, value: float | None) -> None: if value is None: return attr_value = getattr(self, "_" + attr_name) if attr_value is None: setattr(self, "_" + attr_name, value) elif not np.isclose(attr_value, value): raise ValueError( f"Computed value of {attr_name} inconsistent with given value" ) def _set_attribute_values( self, attr_names: Sequence[str], values: Sequence[float | None] ) -> None: for attr_name, value in zip(attr_names, values): self._set_attribute_value(attr_name, value) def __getnewargs__(self) -> tuple[float | None, float | None, float | None]: return self.x, self.y, self.z @staticmethod def _cartesian_to_other( x: float | None, y: float | None, z: float | None ) -> AllCoordsType | None: """Convert a cartesian set of coordinates to values of interest.""" if x is None or y is None or z is None: return None phi = np.arctan2(y, x) rho = np.sqrt(x**2 + y**2) r = np.sqrt(x**2 + y**2 + z**2) if r != 0: z_r_frac = z / r # it it possible that z_r_frac is slightly larger than 1 or # slightly smaller than -1 due to floating point errors. # an example that triggers this is: # x=0, y=0, z=3.729170476738041e-155 if z_r_frac > 1: z_r_frac = 1 elif z_r_frac < -1: z_r_frac = -1 theta = np.arccos(z_r_frac) else: theta = 0 return x, y, z, r, theta, phi, rho @staticmethod def _spherical_to_other( r: float | None, theta: float | None, phi: float | None ) -> AllCoordsType | None: """Convert from spherical to other representations.""" if r is None or theta is None or phi is None: return None z = r * np.cos(theta) x = r * np.sin(theta) * np.cos(phi) y = r * np.sin(theta) * np.sin(phi) rho = np.sqrt(x**2 + y**2) return x, y, z, r, theta, phi, rho @staticmethod def _cylindrical_to_other( phi: float | None, rho: float | None, z: float | None ) -> AllCoordsType | None: """Convert from cylindrical to other representations.""" if phi is None or rho is None or z is None: return None x = rho * np.cos(phi) y = rho * np.sin(phi) r = np.sqrt(rho**2 + z**2) if r != 0: z_r_frac = z / r # it it possible that z_r_frac is slightly larger than 1 or # slightly smaller than -1 due to floating point errors. # an example that triggers this is: # phi=0, rho=0, z=3.729170476738041e-155 if z_r_frac > 1: z_r_frac = 1 elif z_r_frac < -1: z_r_frac = -1 theta = np.arccos(z_r_frac) else: theta = 0 return x, y, z, r, theta, phi, rho def _compute_unknowns(self) -> None: """ Compute all coordinates. To do this we need either the set (x, y, z) to contain no ``None`` values, or the set (r, theta, phi), or the set (rho, phi, z). Given any of these sets, we can recompute the rest. This function will raise an error if there are contradictory inputs (e.g. x=3, y=4, z=0 and rho=6). """ for f in [ lambda: FieldVector._cartesian_to_other(self._x, self._y, self._z), lambda: FieldVector._spherical_to_other(self._r, self._theta, self._phi), lambda: FieldVector._cylindrical_to_other(self._phi, self._rho, self._z), ]: new_values = f() if new_values is not None: # this will return None if any of the # function arguments is None. self._set_attribute_values(FieldVector.attributes, new_values) break
[docs] def copy(self: T, other: T) -> None: """Copy the properties of other vector to yourself.""" for att in FieldVector.attributes: value = getattr(other, "_" + att) setattr(self, "_" + att, value)
[docs] def set_vector(self, **new_values: float) -> None: """ Reset the the values of the vector. Examples: >>> f = FieldVector(x=0, y=2, z=6) >>> f.set_vector(x=9, y=3, z=1) >>> f.set_vector(r=1, theta=30.0, phi=10.0) # The following should raise a value error: # "Can only set vector with a complete value set" >>> f.set_vector(x=9, y=0) # Although mathematically it is possible to compute the complete # vector from the values given, this is too hard to implement with # generality (and not worth it), so the following will raise the # above-mentioned ValueError too. >>> f.set_vector(x=9, y=0, r=3) """ names = sorted(list(new_values.keys())) groups = [["x", "y", "z"], ["phi", "r", "theta"], ["phi", "rho", "z"]] if names not in groups: raise ValueError("Can only set vector with a complete value set") new_vector = self.__class__(**new_values) self.copy(new_vector)
[docs] def set_component(self, **new_values: float) -> None: """ Set a single component of the vector to some new value. It is disallowed for the user to set vector components manually as this can lead to inconsistencies (e.g. x and rho are not independent of each other, setting one has to effect the other). Examples: >>> f = FieldVector(x=2, y=3, z=4) # Since r is part of the set (r, theta, phi) representing # spherical coordinates, setting r means that theta and phi are # kept constant and only r is changed. After changing r, # (x, y, z) values are recomputed, as is the rho coordinate. # Internally we arrange this by setting x, y, z and rho to None # and calling self._compute_unknowns(). >>> f.set_component(r=10) Args: new_values (dict): Keys representing parameter names and values the values to be set. """ if len(new_values) > 1: raise NotImplementedError("Cannot set multiple components at once") items = list(new_values.items()) component_name = items[0][0] if component_name in ["theta", "phi"]: # convert angles to radians value = np.radians(items[0][1]) else: value = items[0][1] setattr(self, "_" + component_name, float(value)) groups = [["x", "y", "z"], ["r", "theta", "phi"], ["phi", "rho", "z"]] for group in groups: if component_name in group: for att in FieldVector.attributes: if att not in group: setattr(self, "_" + att, None) break self._compute_unknowns()
[docs] def get_components(self, *names: str) -> list[float]: """Get field components by name.""" def convert_angle_to_degrees(name: str, value: float) -> float: # Convert all angles to degrees if name in ["theta", "phi"]: return float(np.degrees(value)) else: return value components = [ convert_angle_to_degrees(name, getattr(self, "_" + name)) for name in names ] return components
[docs] def is_equal(self, other: FieldVector) -> bool: """ Returns ``True`` if ``other`` is equivalent to ``self``, ``False`` otherwise. """ for name in ["x", "y", "z"]: self_value = getattr(self, name) other_value = getattr(other, name) if not np.isclose(self_value, other_value): return False return True
def __getitem__(self, component: str) -> float: return self.get_components(component)[0] def __setitem__(self, component: str, value: float) -> None: self.set_component(**{component: value}) def __mul__(self, other: Any) -> FieldVector: if not isinstance(other, (float, int)): return NotImplemented return FieldVector( **{component: self[component] * other for component in "xyz"} ) def __rmul__(self, other: Any) -> FieldVector: if not isinstance(other, (int, float)): return NotImplemented return self * other def __truediv__(self, other: Any) -> FieldVector: if not isinstance(other, (int, float)): return NotImplemented return self * (1.0 / other) def __neg__(self) -> FieldVector: return -1 * self def __add__(self, other: Any) -> FieldVector: if not isinstance(other, FieldVector): return NotImplemented return FieldVector( **{component: self[component] + other[component] for component in "xyz"} ) def __sub__(self, other: Any) -> FieldVector: if not isinstance(other, FieldVector): return NotImplemented return FieldVector( **{component: self[component] - other[component] for component in "xyz"} )
[docs] def norm(self, ord: NormOrder = 2) -> float: """ Returns the norm of this field vector. See ``np.norm`` for the definition of the ord keyword argument. """ assert self.x is not None assert self.y is not None assert self.z is not None return float(np.linalg.norm([self.x, self.y, self.z], ord=ord))
[docs] def distance( self, other: FieldVector, ord: NormOrder = 2, ) -> float: return (self - other).norm(ord=ord)
@property def x(self) -> float | None: return self._x @property def y(self) -> float | None: return self._y @property def z(self) -> float | None: return self._z @property def rho(self) -> float | None: return self._rho @property def theta(self) -> float | None: if self._theta is None: return None return float(np.degrees(self._theta)) @property def r(self) -> float | None: return self._r @property def phi(self) -> float | None: if self._phi is None: return None return float(np.degrees(self._phi)) # Representation Methods #
[docs] def repr_cartesian(self) -> str: return f"FieldVector(x={self.x}, y={self.y}, z={self.z})"
[docs] def repr_spherical(self) -> str: return f"FieldVector(r={self.r}, phi={self.phi}, theta={self.theta})"
[docs] def repr_cylindrical(self) -> str: return f"FieldVector(rho={self.rho}, phi={self.phi}, z={self.z})"
def __repr__(self) -> str: if self.repr_format == "cartesian": return self.repr_cartesian() elif self.repr_format == "spherical": return self.repr_spherical() elif self.repr_format == "cylindrical": return self.repr_cylindrical() else: return super().__repr__() # Homogeneous Coordinates #
[docs] def as_homogeneous(self) -> np.ndarray: return np.array([self.x, self.y, self.z, 1])
[docs] @classmethod def from_homogeneous(cls: type[T], hvec: np.ndarray) -> T: # Homogeneous coordinates define an equivalence relation # [x / s, y / s, z / s, 1] == [x, y, z, s]. # More generally, # [x, y, z, s] == [x', y', z', s'] # iff x / s == x' / s', # y / s == y' / s', and # z / s == z' / s'. # This definition has the consequence that for any w, # w * [x, y, z, s] == [x, y, z, s]. # Thus, we start by rescaling such that s == 1. hvec /= hvec[-1] return cls(x=hvec[0], y=hvec[1], z=hvec[2])