Update type hints and docstrings in plasmapy.formulary (#2543)

* Add type hints & simplify particle usage

* Remove a mypy error to ignore

* Update imports & names in formulary.collisions

* Make imports more specific

* Make imports more specific in formulary

* Update type hints & docstrings for formulary.dielectric

* Regenerate mypy.ini

* Put dielectric tensor output types in __all__ & add docstrings

* Edit changelog: 2543.internal.rst

* Edit changelog: 2543.trivial.rst

changelog/2543.internal.rst

@@ -0,0 +1 @@
+Added and updated type hint annotations across `plasmapy.formulary`.

changelog/2543.trivial.rst

@@ -0,0 +1,3 @@
+Exposed `~plasmapy.formulary.dielectric.StixTensorElements`
+and `~plasmapy.formulary.dielectric.RotatingTensorElements`
+to the public API.

mypy.ini

@@ -189,7 +189,7 @@ disable_error_code = arg-type,attr-defined,misc,no-untyped-call,no-untyped-def,t
 disable_error_code = attr-defined,misc,name-defined,no-untyped-call,no-untyped-def,syntax
 
 [mypy-plasmapy.formulary.collisions.dimensionless]
-disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def
+disable_error_code = misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.collisions.frequencies]
 disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def,union-attr
@@ -201,10 +201,10 @@ disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def
 disable_error_code = attr-defined,no-untyped-def
 
 [mypy-plasmapy.formulary.collisions.lengths]
-disable_error_code = attr-defined,misc,name-defined,no-untyped-call,no-untyped-def,syntax
+disable_error_code = misc,no-untyped-call,no-untyped-def,syntax
 
 [mypy-plasmapy.formulary.collisions.misc]
-disable_error_code = attr-defined,misc,name-defined,no-untyped-call,no-untyped-def,syntax
+disable_error_code = misc,no-untyped-call,no-untyped-def,syntax
 
 [mypy-plasmapy.formulary.collisions.tests.test_coulomb]
 disable_error_code = attr-defined,no-untyped-def
@@ -222,13 +222,13 @@ disable_error_code = attr-defined,no-untyped-def
 disable_error_code = attr-defined,no-untyped-def
 
 [mypy-plasmapy.formulary.densities]
-disable_error_code = attr-defined,misc,no-untyped-call,syntax
+disable_error_code = misc,no-untyped-call,syntax
 
 [mypy-plasmapy.formulary.dielectric]
 disable_error_code = misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.dimensionless]
-disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def
+disable_error_code = misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.distribution]
 disable_error_code = attr-defined,no-untyped-call,no-untyped-def
@@ -243,7 +243,7 @@ disable_error_code = attr-defined,misc,no-any-return,no-untyped-call,union-attr
 disable_error_code = misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.lengths]
-disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def
+disable_error_code = misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.magnetostatics]
 disable_error_code = misc,no-untyped-call,no-untyped-def
@@ -255,16 +255,16 @@ disable_error_code = no-any-return,type-arg
 disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def
 
 [mypy-plasmapy.formulary.quantum]
-disable_error_code = attr-defined,misc,no-untyped-call,no-untyped-def,union-attr
+disable_error_code = misc,no-untyped-call,no-untyped-def,union-attr
 
 [mypy-plasmapy.formulary.radiation]
-disable_error_code = attr-defined,misc,no-any-return,no-untyped-call,type-arg,union-attr
+disable_error_code = misc,no-untyped-call,union-attr
 
 [mypy-plasmapy.formulary.relativity]
-disable_error_code = attr-defined,misc,no-any-return,no-redef,no-untyped-call,no-untyped-def,return-value,union-attr
+disable_error_code = misc,no-any-return,no-redef,no-untyped-call,no-untyped-def,return-value,union-attr
 
 [mypy-plasmapy.formulary.speeds]
-disable_error_code = attr-defined,misc,no-any-return,no-untyped-call,no-untyped-def,syntax,union-attr
+disable_error_code = misc,no-any-return,no-untyped-call,no-untyped-def,syntax,union-attr
 
 [mypy-plasmapy.formulary.tests.test_densities]
 disable_error_code = attr-defined,no-untyped-def

plasmapy/formulary/collisions/dimensionless.py

@@ -176,8 +176,8 @@ def coupling_parameter(
         # using mean charge to get average ion density.
         # If you are running this, you should strongly consider giving
         # a value of z_mean as an argument instead.
-        Z1 = np.abs(particles.charge_number(species[0]))
-        Z2 = np.abs(particles.charge_number(species[1]))
+        Z1 = np.abs(particles.atomic.charge_number(species[0]))
+        Z2 = np.abs(particles.atomic.charge_number(species[1]))
         Z = (Z1 + Z2) / 2
         # getting ion density from electron density
         n_i = n_e / Z

plasmapy/formulary/collisions/frequencies.py

@@ -18,7 +18,7 @@
 from plasmapy import particles
 from plasmapy.formulary.collisions import coulomb, lengths, misc
 from plasmapy.formulary.speeds import thermal_speed
-from plasmapy.particles import ParticleLike
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import deprecated, validate_quantities
 from plasmapy.utils.exceptions import PhysicsError, PlasmaPyFutureWarning
 

plasmapy/formulary/collisions/lengths.py

@@ -8,20 +8,21 @@
 import numpy as np
 from astropy.constants.si import eps0, hbar
 
-from plasmapy import particles
 from plasmapy.formulary.collisions import frequencies, misc
 from plasmapy.formulary.lengths import Debye_length
 from plasmapy.formulary.quantum import Wigner_Seitz_radius
+from plasmapy.particles.decorators import particle_input
+from plasmapy.particles.particle_class import Particle
 from plasmapy.utils.decorators import validate_quantities
 
 
 @validate_quantities(
     T={"can_be_negative": False, "equivalencies": u.temperature_energy()}
 )
-@particles.particle_input
+@particle_input
 def impact_parameter_perp(
     T: u.Quantity[u.K],
-    species: (particles.Particle, particles.Particle),
+    species: (Particle, Particle),
     V: u.Quantity[u.m / u.s] = np.nan * u.m / u.s,
 ) -> u.Quantity[u.m]:
     r"""

plasmapy/formulary/collisions/misc.py

@@ -11,19 +11,19 @@
 import numpy as np
 from astropy.constants.si import e
 
-from plasmapy import particles
 from plasmapy.formulary.collisions import frequencies
 from plasmapy.formulary.speeds import thermal_speed
+from plasmapy.particles.atomic import reduced_mass
+from plasmapy.particles.decorators import particle_input
+from plasmapy.particles.particle_class import Particle
 from plasmapy.utils.decorators import validate_quantities
 from plasmapy.utils.decorators.checks import _check_relativistic
 from plasmapy.utils.exceptions import PhysicsError
 
 
 @validate_quantities(T={"equivalencies": u.temperature_energy()})
-@particles.particle_input
-def _process_inputs(
-    T: u.Quantity[u.K], species: (particles.Particle, particles.Particle), V
-):
+@particle_input
+def _process_inputs(T: u.Quantity[u.K], species: (Particle, Particle), V):
     """
     Helper function for processing inputs to functionality contained
     in `plasmapy.formulary.collisions`.
@@ -35,22 +35,22 @@ def _process_inputs(
     charges = [np.abs(p.charge) for p in species]
 
     # obtaining reduced mass of 2 particle collision system
-    reduced_mass = particles.reduced_mass(*species)
+    reduced_mass_ = reduced_mass(*species)
 
     # getting thermal velocity of system if no velocity is given
-    V = _replace_nan_velocity_with_thermal_velocity(V, T, reduced_mass)
+    V = _replace_nan_velocity_with_thermal_velocity(V, T, reduced_mass_)
 
     _check_relativistic(V, "V")
 
-    return T, masses, charges, reduced_mass, V
+    return T, masses, charges, reduced_mass_, V
 
 
 # TODO: Remove redundant mass parameter
 def _replace_nan_velocity_with_thermal_velocity(
     V,
     T,
     m,
-    species=particles.Particle("e-"),  # noqa: B008
+    species=Particle("e-"),  # noqa: B008
 ):
     """
     Get thermal velocity of system if no velocity is given, for a given
@@ -203,9 +203,9 @@ def mobility(
     # we do this after collision_frequency since collision_frequency
     # already has a _process_inputs check and we are doing this just
     # to recover the charges, mass, etc.
-    T, masses, charges, reduced_mass, V = _process_inputs(T=T, species=species, V=V)
+    T, masses, charges, reduced_mass_, V = _process_inputs(T=T, species=species, V=V)
     z_val = (charges[0] + charges[1]) / 2 if np.isnan(z_mean) else z_mean * e
-    return z_val / (reduced_mass * freq)
+    return z_val / (reduced_mass_ * freq)
 
 
 @validate_quantities(
@@ -328,9 +328,9 @@ def Spitzer_resistivity(
         T=T, n=n, species=species, z_mean=z_mean, V=V, method=method
     )
     # fetching additional parameters
-    T, masses, charges, reduced_mass, V = _process_inputs(T=T, species=species, V=V)
+    T, masses, charges, reduced_mass_, V = _process_inputs(T=T, species=species, V=V)
     return (
-        freq * reduced_mass / (n * charges[0] * charges[1])
+        freq * reduced_mass_ / (n * charges[0] * charges[1])
         if np.isnan(z_mean)
-        else freq * reduced_mass / (n * (z_mean * e) ** 2)
+        else freq * reduced_mass_ / (n * (z_mean * e) ** 2)
     )

plasmapy/formulary/densities.py

@@ -9,7 +9,7 @@
 import numpy as np
 from astropy.constants.si import e, eps0, m_e
 
-from plasmapy.particles import Particle, ParticleLike
+from plasmapy.particles.particle_class import Particle, ParticleLike
 from plasmapy.utils.decorators import validate_quantities
 
 __all__ += __aliases__

plasmapy/formulary/dielectric.py

@@ -3,17 +3,22 @@
     "cold_plasma_permittivity_SDP",
     "cold_plasma_permittivity_LRP",
     "permittivity_1D_Maxwellian",
+    "RotatingTensorElements",
+    "StixTensorElements",
 ]
 __lite_funcs__ = ["permittivity_1D_Maxwellian_lite"]
 
 from collections import namedtuple
+from collections.abc import Sequence
 
 import astropy.units as u
 import numpy as np
 
 from plasmapy.dispersion.dispersion_functions import plasma_dispersion_func_deriv
 from plasmapy.formulary.frequencies import gyrofrequency, plasma_frequency
 from plasmapy.formulary.speeds import thermal_speed
+from plasmapy.particles.particle_class import ParticleLike
+from plasmapy.particles.particle_collections import ParticleListLike
 from plasmapy.utils.decorators import (
     bind_lite_func,
     preserve_signature,
@@ -27,15 +32,22 @@
 """
 
 StixTensorElements = namedtuple("StixTensorElements", ["sum", "difference", "plasma"])
+"""Output type for `~plasmapy.formulary.dielectric.cold_plasma_permittivity_SDP`."""
+
+
 RotatingTensorElements = namedtuple(
     "RotatingTensorElements", ["left", "right", "plasma"]
 )
+"""Output type for `~plasmapy.formulary.dielectric.cold_plasma_permittivity_LRP`."""
 
 
 @validate_quantities(B={"can_be_negative": False}, omega={"can_be_negative": False})
 def cold_plasma_permittivity_SDP(
-    B: u.Quantity[u.T], species, n, omega: u.Quantity[u.rad / u.s]
-):
+    B: u.Quantity[u.T],
+    species: ParticleListLike,
+    n: Sequence[u.Quantity[u.m**-3]] | u.Quantity[u.m**-3],
+    omega: u.Quantity[u.rad / u.s],
+) -> StixTensorElements:
     r"""
     Magnetized cold plasma dielectric permittivity tensor elements.
 
@@ -49,7 +61,7 @@ def cold_plasma_permittivity_SDP(
     B : `~astropy.units.Quantity`
         Magnetic field magnitude in units convertible to tesla.
 
-    species : `list` of `str`
+    species : |particle-list-like|
         List of the plasma particle species,
         e.g.: ``['e-', 'D+']`` or ``['e-', 'D+', 'He+']``.
 
@@ -63,13 +75,13 @@ def cold_plasma_permittivity_SDP(
     Returns
     -------
     sum : `~astropy.units.Quantity`
-        S ("Sum") dielectric tensor element.
+        The "sum" dielectric tensor element, :math:`S`.
 
     difference : `~astropy.units.Quantity`
-        D ("Difference") dielectric tensor element.
+        The "difference" dielectric tensor element, :math:`D`.
 
     plasma : `~astropy.units.Quantity`
-        P ("Plasma") dielectric tensor element.
+        The "plasma" dielectric tensor element, :math:`P`.
 
     Notes
     -----
@@ -129,8 +141,11 @@ def cold_plasma_permittivity_SDP(
 
 @validate_quantities(B={"can_be_negative": False}, omega={"can_be_negative": False})
 def cold_plasma_permittivity_LRP(
-    B: u.Quantity[u.T], species, n, omega: u.Quantity[u.rad / u.s]
-):
+    B: u.Quantity[u.T],
+    species: ParticleListLike,
+    n: list[u.Quantity[u.m**-3]] | u.Quantity[u.m**-3],
+    omega: u.Quantity[u.rad / u.s],
+) -> RotatingTensorElements:
     r"""
     Magnetized cold plasma dielectric permittivity tensor elements.
 
@@ -145,11 +160,11 @@ def cold_plasma_permittivity_LRP(
     B : `~astropy.units.Quantity`
         Magnetic field magnitude in units convertible to tesla.
 
-    species : `list` of `str`
+    species : (k,) |particle-list-like|
         The plasma particle species (e.g.: ``['e-', 'D+']`` or
         ``['e-', 'D+', 'He+']``.
 
-    n : `list` of `~astropy.units.Quantity`
+    n : (k,) `list` of `~astropy.units.Quantity`
         `list` of species density in units convertible to per cubic meter.
         The order of the species densities should follow species.
 
@@ -288,7 +303,7 @@ def permittivity_1D_Maxwellian(
     kWave: u.Quantity[u.rad / u.m],
     T: u.Quantity[u.K],
     n: u.Quantity[u.m**-3],
-    particle,
+    particle: ParticleLike,
     z_mean: float | None = None,
 ) -> u.Quantity[u.dimensionless_unscaled]:
     r"""
@@ -318,7 +333,7 @@ def permittivity_1D_Maxwellian(
         The plasma density — this can be either the electron or the ion
         density, but should be consistent with temperature and particle.
 
-    particle : `str`
+    particle : |particle-like|
         The plasma particle species.
 
     z_mean : `float`

plasmapy/formulary/dimensionless.py

@@ -26,7 +26,8 @@
 
 from plasmapy.formulary import frequencies, lengths, misc, speeds
 from plasmapy.formulary.quantum import quantum_theta
-from plasmapy.particles import Particle, ParticleLike, particle_input
+from plasmapy.particles.decorators import particle_input
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import validate_quantities
 
 __all__ += __aliases__
@@ -118,9 +119,9 @@ def Hall_parameter(
     B: u.Quantity[u.T],
     ion: ParticleLike,
     particle: ParticleLike,
-    coulomb_log=None,
-    V=None,
-    coulomb_log_method="classical",
+    coulomb_log: float | None = None,
+    V: u.Quantity[u.m / u.s] | None = None,
+    coulomb_log_method: str = "classical",
 ):
     r"""
     Calculate the ``particle`` Hall parameter for a plasma.
@@ -215,7 +216,7 @@ def Hall_parameter(
 
     gyro_frequency = frequencies.gyrofrequency(B, particle)
     gyro_frequency = gyro_frequency / u.radian
-    if Particle(particle).symbol == "e-":
+    if particle == "e-":
         coll_rate = fundamental_electron_collision_freq(
             T, n, ion, coulomb_log, V, coulomb_log_method=coulomb_log_method
         )

plasmapy/formulary/frequencies.py

@@ -16,8 +16,9 @@
 from numba import njit
 
 from plasmapy import particles
-from plasmapy.particles import ParticleLike, particle_input
+from plasmapy.particles.decorators import particle_input
 from plasmapy.particles.exceptions import InvalidParticleError
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import (
     angular_freq_to_hz,
     bind_lite_func,

plasmapy/formulary/lengths.py

@@ -11,7 +11,8 @@
 
 from plasmapy.formulary import frequencies, speeds
 from plasmapy.formulary.relativity import RelativisticBody
-from plasmapy.particles import ParticleLike, particle_input
+from plasmapy.particles.decorators import particle_input
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import validate_quantities
 
 __all__ += __aliases__

plasmapy/formulary/misc.py

@@ -12,7 +12,7 @@
 from astropy.constants.si import e, k_B, mu0
 
 from plasmapy import particles
-from plasmapy.particles import ParticleLike
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import validate_quantities
 
 __all__ += __aliases__

plasmapy/formulary/quantum.py

@@ -21,7 +21,8 @@
 
 from plasmapy.formulary import mathematics
 from plasmapy.formulary.relativity import Lorentz_factor
-from plasmapy.particles import ParticleLike, particle_input
+from plasmapy.particles.decorators import particle_input
+from plasmapy.particles.particle_class import ParticleLike
 from plasmapy.utils.decorators import validate_quantities
 from plasmapy.utils.exceptions import RelativityError