Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update type hints and docstrings in plasmapy.formulary #2543

Merged
merged 10 commits into from
Feb 28, 2024
1 change: 1 addition & 0 deletions changelog/2543.internal.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Added and updated type hint annotations across `plasmapy.formulary`.
3 changes: 3 additions & 0 deletions changelog/2543.trivial.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Exposed `~plasmapy.formulary.dielectric.StixTensorElements`
and `~plasmapy.formulary.dielectric.RotatingTensorElements`
to the public API.
20 changes: 10 additions & 10 deletions mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
4 changes: 2 additions & 2 deletions plasmapy/formulary/collisions/dimensionless.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion plasmapy/formulary/collisions/frequencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There was a mypy error about how plasmapy.particles doesn't explicitly export ParticleLike. Hence, I'm making a bunch of these import statements a lot more specific.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think I remember this issue

from plasmapy.utils.decorators import deprecated, validate_quantities
from plasmapy.utils.exceptions import PhysicsError, PlasmaPyFutureWarning

Expand Down
7 changes: 4 additions & 3 deletions plasmapy/formulary/collisions/lengths.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Comment on lines -21 to +25
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be more consistent with the import conventions we use elsewhere.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I love it when a review is just basically import statements, it makes the review a lot easier 😆

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It makes the pull request a lot easier too!

V: u.Quantity[u.m / u.s] = np.nan * u.m / u.s,
) -> u.Quantity[u.m]:
r"""
Expand Down
28 changes: 14 additions & 14 deletions plasmapy/formulary/collisions/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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)
)
2 changes: 1 addition & 1 deletion plasmapy/formulary/densities.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
Expand Down
39 changes: 27 additions & 12 deletions plasmapy/formulary/dielectric.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,22 @@
"cold_plasma_permittivity_SDP",
"cold_plasma_permittivity_LRP",
"permittivity_1D_Maxwellian",
"RotatingTensorElements",
"StixTensorElements",
Comment on lines +6 to +7
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I needed to expose these to the public API so that they could be used as return annotations without causing a "missing reference" error when doing the doc build.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you explain this when we next talk please?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be happy to, with the caveat that I only understand 83% of what's happening here!

]
__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,
Expand All @@ -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.

Expand All @@ -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+']``.

Expand All @@ -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
-----
Expand Down Expand Up @@ -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.

Expand All @@ -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.

Expand Down Expand Up @@ -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"""
Expand Down Expand Up @@ -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`
Expand Down
11 changes: 6 additions & 5 deletions plasmapy/formulary/dimensionless.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
)
Expand Down
3 changes: 2 additions & 1 deletion plasmapy/formulary/frequencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion plasmapy/formulary/lengths.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
Expand Down
2 changes: 1 addition & 1 deletion plasmapy/formulary/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
Expand Down
3 changes: 2 additions & 1 deletion plasmapy/formulary/quantum.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down