diff --git a/.gitignore b/.gitignore index 11edeed..9e6c649 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ docs/source/examples docs/source/reference/generated +docs/source/sg_execution_times.rst led / optimized / DLL files __pycache__/ diff --git a/LICENSE b/LICENSE index a43592e..85d18c2 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2018-2020 Greg Lucas +Copyright (c) 2018 Greg Lucas Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/docs/source/conf.py b/docs/source/conf.py index 1762182..2049642 100755 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -19,10 +19,12 @@ # import os import sys + import pysecs -sys.path.append(os.path.abspath('..')) -sys.path.append('..') + +sys.path.append(os.path.abspath("..")) +sys.path.append("..") # -- General configuration --------------------------------------------- @@ -33,37 +35,36 @@ # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones. -extensions = ['sphinx.ext.autodoc', - 'sphinx.ext.autosummary', - 'sphinx.ext.githubpages', # Helpful for publishing to gh-pages - 'sphinx.ext.intersphinx', - 'sphinx.ext.napoleon', - 'sphinx_gallery.gen_gallery', - ] +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.githubpages", # Helpful for publishing to gh-pages + "sphinx.ext.intersphinx", + "sphinx.ext.napoleon", + "sphinx_gallery.gen_gallery", +] autosummary_generate = True napoleon_google_docstring = False -intersphinx_mapping = { - 'numpy': ('https://docs.scipy.org/doc/numpy/', None) -} +intersphinx_mapping = {"numpy": ("https://numpy.org/doc/stable/", None)} # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = '.rst' +source_suffix = ".rst" # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = 'pysecs' -copyright = "2018-2021, Greg Lucas" +project = "pysecs" +copyright = "2018-, Greg Lucas" author = "Greg Lucas" # The version info for the project you're documenting, acts as replacement @@ -75,20 +76,13 @@ # The full version, including alpha/beta/rc tags. release = pysecs.__version__ -# The language for content autogenerated by Sphinx. Refer to documentation -# for a list of supported languages. -# -# This is also used if you do content translation via gettext catalogs. -# Usually you set "language" from the command line for these cases. -language = None - # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'lovelace' +pygments_style = "lovelace" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False @@ -99,12 +93,12 @@ # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -html_theme = 'pydata_sphinx_theme' +html_theme = "pydata_sphinx_theme" html_logo = "_static/pysecs.png" html_theme_options = { - "github_url": "https://github.com/greglucas/pysecs", + "github_url": "https://github.com/greglucas/pysecs", } # Add any paths that contain custom static files (such as style sheets) here, @@ -116,16 +110,16 @@ # Sphinx gallery sphinx_gallery_conf = { - 'examples_dirs': '../../examples', # path to example scripts - 'gallery_dirs': 'examples', # path for where to save generated output - 'matplotlib_animations': True, + "examples_dirs": "../../examples", # path to example scripts + "gallery_dirs": "examples", # path for where to save generated output + "matplotlib_animations": True, } # -- Options for HTMLHelp output --------------------------------------- # Output file base name for HTML help builder. -htmlhelp_basename = 'pysecsdoc' +htmlhelp_basename = "pysecsdoc" # -- Options for LaTeX output ------------------------------------------ @@ -134,15 +128,12 @@ # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -152,9 +143,7 @@ # (source start file, target name, title, author, documentclass # [howto, manual, or own class]). latex_documents = [ - (master_doc, 'pysecs.tex', - 'pysecs Documentation', - 'Greg Lucas', 'manual'), + (master_doc, "pysecs.tex", "pysecs Documentation", "Greg Lucas", "manual"), ] @@ -162,11 +151,7 @@ # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'pysecs', - 'pysecs Documentation', - [author], 1) -] +man_pages = [(master_doc, "pysecs", "pysecs Documentation", [author], 1)] # -- Options for Texinfo output ---------------------------------------- @@ -175,10 +160,13 @@ # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'pysecs', - 'pysecs Documentation', - author, - 'pysecs', - 'Spherical Elementary Current Systems in Python.', - 'Miscellaneous'), + ( + master_doc, + "pysecs", + "pysecs Documentation", + author, + "pysecs", + "Spherical Elementary Current Systems in Python.", + "Miscellaneous", + ), ] diff --git a/examples/plot_B_divergence_free.py b/examples/plot_B_divergence_free.py index 51fb244..7c375eb 100644 --- a/examples/plot_B_divergence_free.py +++ b/examples/plot_B_divergence_free.py @@ -6,15 +6,18 @@ an SECS pole placed at the north pole. Figure 2 from Amm and Viljanen 1999. """ + import matplotlib.pyplot as plt import numpy as np + from pysecs import SECS + # Radius of Earth R_E = 6378e3 # Pole of the current system at the North Pole -sec_loc = np.array([90., 0., R_E + 100e3]) +sec_loc = np.array([90.0, 0.0, R_E + 100e3]) # Set up the system with a single SEC system_df = SECS(sec_df_loc=sec_loc) @@ -33,27 +36,27 @@ pred_loc = np.zeros(shape=(N, 3)) angles = np.linspace(0, 45, N) -pred_loc[:, 0] = 90-angles +pred_loc[:, 0] = 90 - angles pred_loc[:, 2] = R_E B_pred = system_df.predict(pred_loc=pred_loc) # B_theta == -Bx == -B_pred[:,0] # Convert to nT (1e9) -B_theta = -B_pred[:, 0]*1e9 +B_theta = -B_pred[:, 0] * 1e9 # B_r == -Bz = -B_pred[:,2] -B_r = -B_pred[:, 2]*1e9 +B_r = -B_pred[:, 2] * 1e9 fig, ax = plt.subplots() -ax.plot(angles, B_theta, c='b', label=r'B$_{\theta}$') -ax.plot(angles, B_r, c='r', label=r'B$_r$') -ax.legend(loc='upper right') +ax.plot(angles, B_theta, c="b", label=r"B$_{\theta}$") +ax.plot(angles, B_r, c="r", label=r"B$_r$") +ax.legend(loc="upper right") ax.set_xlim(angles[0], angles[-1]) ax.set_ylim(-5, 10) -ax.set_xlabel('Angle from pole (deg)') -ax.set_ylabel('B (nT)') -ax.axhline(0., c='k', alpha=0.5, linewidth=0.5) -ax.set_title('Amm and Viljanen (1999) Figure 2') +ax.set_xlabel("Angle from pole (deg)") +ax.set_ylabel("B (nT)") +ax.axhline(0.0, c="k", alpha=0.5, linewidth=0.5) +ax.set_title("Amm and Viljanen (1999) Figure 2") plt.show() diff --git a/examples/plot_J_surface.py b/examples/plot_J_surface.py index ad07ef0..f8ad9fa 100644 --- a/examples/plot_J_surface.py +++ b/examples/plot_J_surface.py @@ -5,16 +5,19 @@ This example demonstrates the resulting current surfaces from the curl-free and divergence-free SEC poles. """ + import matplotlib.pyplot as plt -from matplotlib.colors import LogNorm import numpy as np +from matplotlib.colors import LogNorm + from pysecs import SECS + # Radius of Earth R_E = 6378e3 # Pole of the current system at the North Pole -sec_loc = np.array([90., 0., R_E + 100e3]) +sec_loc = np.array([90.0, 0.0, R_E + 100e3]) # Set up the system with a single SEC system_df = SECS(sec_df_loc=sec_loc) @@ -39,8 +42,8 @@ lats = np.arange(lat_min, lat_max, delta) lons = np.arange(lon_min, lon_max, delta) -lats2 = np.arange(lat_min-delta/2, lat_max+delta/2, delta) -lons2 = np.arange(lon_min-delta/2, lon_max+delta/2, delta) +lats2 = np.arange(lat_min - delta / 2, lat_max + delta / 2, delta) +lons2 = np.arange(lon_min - delta / 2, lon_max + delta / 2, delta) nlat = len(lats) nlon = len(lons) @@ -48,7 +51,7 @@ xx, yy = np.meshgrid(lons, lats) xx2, yy2 = np.meshgrid(lons2, lats2) -points = np.zeros((nlat*nlon, 3)) +points = np.zeros((nlat * nlon, 3)) points[:, 0] = yy.ravel() points[:, 1] = xx.ravel() points[:, 2] = R_E @@ -70,33 +73,38 @@ J_cf_theta = np.linalg.norm(J_cf[:, :2], axis=1).reshape(nlat, nlon) J_cf_r = J_cf[:, 2].reshape(nlat, nlon) -fig, (ax_cf, ax_df) = plt.subplots(figsize=(8, 3), ncols=2, - constrained_layout=True) -fig.suptitle('Amm and Viljanen (1999) Figure 1') +fig, (ax_cf, ax_df) = plt.subplots(figsize=(8, 3), ncols=2, constrained_layout=True) +fig.suptitle("Amm and Viljanen (1999) Figure 1") # Limit quiver to only a few points nstep = 30 norm = LogNorm(vmin=1e-5, vmax=1e-2) cax = ax_cf.pcolormesh(xx2, yy2, J_cf_theta, norm=norm) -ax_cf.quiver(xx[nstep//2::nstep, nstep//2::nstep*2], - yy[nstep//2::nstep, nstep//2::nstep*2], - J_cf_y[nstep//2::nstep, nstep//2::nstep*2], - J_cf_x[nstep//2::nstep, nstep//2::nstep*2], - color='w', zorder=9) -ax_cf.set_ylabel('Latitude (deg)') -ax_cf.set_xlabel('Longitude (deg)') +ax_cf.quiver( + xx[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + yy[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + J_cf_y[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + J_cf_x[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + color="w", + zorder=9, +) +ax_cf.set_ylabel("Latitude (deg)") +ax_cf.set_xlabel("Longitude (deg)") cax = ax_df.pcolormesh(xx2, yy2, J_df_phi, norm=norm) -ax_df.quiver(xx[nstep//2::nstep, nstep//2::nstep*2], - yy[nstep//2::nstep, nstep//2::nstep*2], - J_df_y[nstep//2::nstep, nstep//2::nstep*2], - J_df_x[nstep//2::nstep, nstep//2::nstep*2], - color='w', zorder=9) -ax_df.set_ylabel('Latitude (deg)') -ax_df.set_xlabel('Longitude (deg)') - -ax_cf.set_title('Curl Free') -ax_df.set_title('Divergence Free') +ax_df.quiver( + xx[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + yy[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + J_df_y[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + J_df_x[nstep // 2 :: nstep, nstep // 2 :: nstep * 2], + color="w", + zorder=9, +) +ax_df.set_ylabel("Latitude (deg)") +ax_df.set_xlabel("Longitude (deg)") + +ax_cf.set_title("Curl Free") +ax_df.set_title("Divergence Free") plt.show() diff --git a/examples/plot_animation.py b/examples/plot_animation.py index 75ca977..3ea4644 100644 --- a/examples/plot_animation.py +++ b/examples/plot_animation.py @@ -7,9 +7,10 @@ to make predictions on a separate grid and compare the results. """ + import matplotlib.pyplot as plt -from matplotlib.animation import FuncAnimation import numpy as np +from matplotlib.animation import FuncAnimation from pysecs import SECS @@ -17,33 +18,29 @@ R_earth = 6371e3 # specify the SECS grid -lat, lon, r = np.meshgrid(np.linspace(-20, 20, 30), - np.linspace(-20, 20, 30), - R_earth + 110000, indexing='ij') -secs_lat_lon_r = np.hstack((lat.reshape(-1, 1), - lon.reshape(-1, 1), - r.reshape(-1, 1))) +lat, lon, r = np.meshgrid( + np.linspace(-20, 20, 30), np.linspace(-20, 20, 30), R_earth + 110000, indexing="ij" +) +secs_lat_lon_r = np.hstack((lat.reshape(-1, 1), lon.reshape(-1, 1), r.reshape(-1, 1))) # Set up the class secs = SECS(sec_df_loc=secs_lat_lon_r) # Make a grid of input observations spanning # (-10, 10) in latitutde and longitude -lat, lon, r = np.meshgrid(np.linspace(-10, 10, 11), - np.linspace(-10, 10, 11), - R_earth, indexing='ij') +lat, lon, r = np.meshgrid( + np.linspace(-10, 10, 11), np.linspace(-10, 10, 11), R_earth, indexing="ij" +) obs_lat = lat[..., 0] obs_lon = lon[..., 0] -obs_lat_lon_r = np.hstack((lat.reshape(-1, 1), - lon.reshape(-1, 1), - r.reshape(-1, 1))) +obs_lat_lon_r = np.hstack((lat.reshape(-1, 1), lon.reshape(-1, 1), r.reshape(-1, 1))) nobs = len(obs_lat_lon_r) # Create the synthetic magnetic field data as a function # of time -ts = np.linspace(0, 2*np.pi) -bx = 5*np.cos(ts) -by = 5*np.sin(ts) +ts = np.linspace(0, 2 * np.pi) +bx = 5 * np.cos(ts) +by = 5 * np.sin(ts) bz = ts # ntimes x 3 B_obs = np.column_stack([bx, by, bz]) @@ -54,8 +51,8 @@ B_obs = np.repeat(B_obs[:, np.newaxis, :], nobs, axis=1) # Make it more interesting and add a sin wave in spatial # coordinates too -B_obs[:, :, 0] *= 2*np.sin(np.deg2rad(obs_lat_lon_r[:, 0])) -B_obs[:, :, 1] *= 2*np.sin(np.deg2rad(obs_lat_lon_r[:, 1])) +B_obs[:, :, 0] *= 2 * np.sin(np.deg2rad(obs_lat_lon_r[:, 0])) +B_obs[:, :, 1] *= 2 * np.sin(np.deg2rad(obs_lat_lon_r[:, 1])) B_std = np.ones(B_obs.shape) # Ignore the Z component @@ -69,52 +66,101 @@ # Create prediction points # Extend it a little beyond the observation points (-11, 11) -lat, lon, r = np.meshgrid(np.linspace(-11, 11, 11), - np.linspace(-11, 11, 11), - R_earth, indexing='ij') -pred_lat_lon_r = np.hstack((lat.reshape(-1, 1), - lon.reshape(-1, 1), - r.reshape(-1, 1))) +lat, lon, r = np.meshgrid( + np.linspace(-11, 11, 11), np.linspace(-11, 11, 11), R_earth, indexing="ij" +) +pred_lat_lon_r = np.hstack((lat.reshape(-1, 1), lon.reshape(-1, 1), r.reshape(-1, 1))) # Call the prediction function B_pred = secs.predict(pred_lat_lon_r) # Now set up the plots -fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2, - sharex=True, sharey=True) +fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2, sharex=True, sharey=True) vmin, vmax = -5, 5 -cmap = plt.get_cmap('RdBu_r') +cmap = plt.get_cmap("RdBu_r") t = 10 -mesh1 = ax1.pcolormesh(obs_lon, obs_lat, B_obs[t, :, 0].reshape(obs_lon.shape), - vmin=vmin, vmax=vmax, cmap=cmap, shading='auto') -mesh2 = ax2.pcolormesh(obs_lon, obs_lat, B_obs[t, :, 1].reshape(obs_lon.shape), - vmin=vmin, vmax=vmax, cmap=cmap, shading='auto') +mesh1 = ax1.pcolormesh( + obs_lon, + obs_lat, + B_obs[t, :, 0].reshape(obs_lon.shape), + vmin=vmin, + vmax=vmax, + cmap=cmap, + shading="auto", +) +mesh2 = ax2.pcolormesh( + obs_lon, + obs_lat, + B_obs[t, :, 1].reshape(obs_lon.shape), + vmin=vmin, + vmax=vmax, + cmap=cmap, + shading="auto", +) ax1.set_ylabel("Observations") ax1.set_title("B$_{X}$") ax2.set_title("B$_{Y}$") qscale = 1 -q1 = ax1.quiver(obs_lat_lon_r[:, 1], obs_lat_lon_r[:, 0], - B_obs[t, :, 1], B_obs[t, :, 0], - angles='xy', scale_units='xy', scale=qscale) -q2 = ax2.quiver(obs_lat_lon_r[:, 1], obs_lat_lon_r[:, 0], - B_obs[t, :, 1], B_obs[t, :, 0], - angles='xy', scale_units='xy', scale=qscale) +q1 = ax1.quiver( + obs_lat_lon_r[:, 1], + obs_lat_lon_r[:, 0], + B_obs[t, :, 1], + B_obs[t, :, 0], + angles="xy", + scale_units="xy", + scale=qscale, +) +q2 = ax2.quiver( + obs_lat_lon_r[:, 1], + obs_lat_lon_r[:, 0], + B_obs[t, :, 1], + B_obs[t, :, 0], + angles="xy", + scale_units="xy", + scale=qscale, +) lon = lon[..., 0] lat = lat[..., 0] -mesh3 = ax3.pcolormesh(lon, lat, B_pred[t, :, 0].reshape(lon.shape), - vmin=vmin, vmax=vmax, cmap=cmap, shading='auto') -mesh4 = ax4.pcolormesh(lon, lat, B_pred[t, :, 1].reshape(lon.shape), - vmin=vmin, vmax=vmax, cmap=cmap, shading='auto') - -q3 = ax3.quiver(pred_lat_lon_r[:, 1], pred_lat_lon_r[:, 0], - B_pred[t, :, 1], B_pred[t, :, 0], - angles='xy', scale_units='xy', scale=qscale) -q4 = ax4.quiver(pred_lat_lon_r[:, 1], pred_lat_lon_r[:, 0], - B_pred[t, :, 1], B_pred[t, :, 0], - angles='xy', scale_units='xy', scale=qscale) +mesh3 = ax3.pcolormesh( + lon, + lat, + B_pred[t, :, 0].reshape(lon.shape), + vmin=vmin, + vmax=vmax, + cmap=cmap, + shading="auto", +) +mesh4 = ax4.pcolormesh( + lon, + lat, + B_pred[t, :, 1].reshape(lon.shape), + vmin=vmin, + vmax=vmax, + cmap=cmap, + shading="auto", +) + +q3 = ax3.quiver( + pred_lat_lon_r[:, 1], + pred_lat_lon_r[:, 0], + B_pred[t, :, 1], + B_pred[t, :, 0], + angles="xy", + scale_units="xy", + scale=qscale, +) +q4 = ax4.quiver( + pred_lat_lon_r[:, 1], + pred_lat_lon_r[:, 0], + B_pred[t, :, 1], + B_pred[t, :, 0], + angles="xy", + scale_units="xy", + scale=qscale, +) ax3.set_ylabel("Predictions") ax3.set_title("B$_{X}$") ax4.set_title("B$_{Y}$") @@ -134,7 +180,6 @@ def update_axes(t): q4.set_UVC(B_pred[t, :, 1], B_pred[t, :, 0]) -ani = FuncAnimation(fig, update_axes, frames=range(ntimes), - interval=50) +ani = FuncAnimation(fig, update_axes, frames=range(ntimes), interval=50) plt.show() diff --git a/pyproject.toml b/pyproject.toml index c99b919..3ea4651 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,7 @@ testpaths = [ ] addopts = [ "--import-mode=importlib", + "--cov", ] [tool.ruff] @@ -77,6 +78,8 @@ lint.ignore = ["B028", "D203", "D212", "N802", "N803", "N806", "PLR0913", "PLR20 [tool.ruff.lint.per-file-ignores] "tests/*" = ["ANN", "D", "S"] +"examples/*" = ["ANN", "D", "S"] +"docs/source/*" = ["ANN", "D", "S"] ".github/*" = ["S"] [tool.ruff.lint.pydocstyle] diff --git a/pysecs/__init__.py b/pysecs/__init__.py index 92a71c7..def1710 100644 --- a/pysecs/__init__.py +++ b/pysecs/__init__.py @@ -1,3 +1,15 @@ -from .secs import * -__all__ = ["SECS"] -__version__ = '0.2.0' +"""Spherical Elementary Current System (SECS) package. + +This package provides a Python implementation of the Spherical Elementary Current +System (SECS) model. It can be used to interpolate between observations +of the magnetic field at Earth's surface. This can be the case when the magnetic +field is assumed to arise from a current sheet in the ionosphere or within the +surface of the Earth. +""" + +import importlib.metadata + +from pysecs.secs import * # noqa + + +__version__ = importlib.metadata.version("pysecs") diff --git a/pysecs/secs.py b/pysecs/secs.py index 704a5d0..b2de3b6 100644 --- a/pysecs/secs.py +++ b/pysecs/secs.py @@ -1,3 +1,8 @@ +"""Spherical Elementary Current System (SECS) module. + +Calculate magnetic field transfer functions and fit a system (SECS) to observations. +""" + import numpy as np @@ -9,69 +14,68 @@ class SECS: Parameters ---------- - - sec_df_loc : ndarray (nsec x 3 [lat, lon, r]) + sec_df_loc : ndarray (nsec, 3 [lat, lon, r]) The latitude, longiutde, and radius of the divergence free (df) SEC locations. - sec_cf_loc : ndarray (nsec x 3 [lat, lon, r]) + sec_cf_loc : ndarray (nsec, 3 [lat, lon, r]) The latitude, longiutde, and radius of the curl free (cf) SEC locations. References ---------- - .. [1] Amm, O., and A. Viljanen. "Ionospheric disturbance magnetic field continuation - from the ground to the ionosphere using spherical elementary current systems." - Earth, Planets and Space 51.6 (1999): 431-440. doi:10.1186/BF03352247 + .. [1] Amm, O., and A. Viljanen. "Ionospheric disturbance magnetic field + continuation from the ground to the ionosphere using spherical elementary + current systems." Earth, Planets and Space 51.6 (1999): 431-440. + doi:10.1186/BF03352247 """ - def __init__(self, sec_df_loc=None, sec_cf_loc=None): - + def __init__( + self, sec_df_loc: np.ndarray | None = None, sec_cf_loc: np.ndarray | None = None + ) -> None: if sec_df_loc is None and sec_cf_loc is None: - raise ValueError("Must initialize the object with SEC locations") - - self.sec_df_loc = sec_df_loc - self.sec_cf_loc = sec_cf_loc - - if self.sec_df_loc is not None: - self.sec_df_loc = np.asarray(sec_df_loc) - if self.sec_df_loc.shape[-1] != 3: - raise ValueError("SEC DF locations must have 3 columns (lat, lon, r)") - if self.sec_df_loc.ndim == 1: - # Add an empty dimension if only one SEC location is passed in - self.sec_df_loc = self.sec_df_loc[np.newaxis, ...] - - if self.sec_cf_loc is not None: - self.sec_cf_loc = np.asarray(sec_cf_loc) - if self.sec_cf_loc.shape[-1] != 3: - raise ValueError("SEC CF locations must have 3 columns (lat, lon, r)") - if self.sec_cf_loc.ndim == 1: - # Add an empty dimension if only one SEC location is passed in - self.sec_cf_loc = self.sec_cf_loc[np.newaxis, ...] + raise ValueError("Must specify at least one of sec_df_loc of sec_cf_loc") + + if sec_df_loc is None: + sec_df_loc = np.empty((0, 3)) + # Convert to array if not already and + # add an empty dimension if only one SEC location is passed in + self.sec_df_loc = np.atleast_2d(sec_df_loc) + if self.sec_df_loc.shape[-1] != 3: + raise ValueError("SEC DF locations must have 3 columns (lat, lon, r)") + + if sec_cf_loc is None: + sec_cf_loc = np.empty((0, 3)) + # Convert to array if not already and + # add an empty dimension if only one SEC location is passed in + self.sec_cf_loc = np.atleast_2d(sec_cf_loc) + if self.sec_cf_loc.shape[-1] != 3: + raise ValueError("SEC CF locations must have 3 columns (lat, lon, r)") # Storage of the scaling factors - self.sec_amps = None - self.sec_amps_var = None + self.sec_amps = np.empty((0, self.nsec)) + self.sec_amps_var = np.empty((0, self.nsec)) @property - def has_df(self): + def has_df(self) -> bool: """Whether this system has any divergence free currents.""" - return self.sec_df_loc is not None + return len(self.sec_df_loc) > 0 @property - def has_cf(self): + def has_cf(self) -> bool: """Whether this system has any curl free currents.""" - return self.sec_cf_loc is not None + return len(self.sec_cf_loc) > 0 @property - def nsec(self): + def nsec(self) -> int: """The number of elementary currents in this system.""" - nsec = 0 - if self.has_df: - nsec += len(self.sec_df_loc) - if self.has_cf: - nsec += len(self.sec_cf_loc) - return nsec - - def fit(self, obs_loc, obs_B, obs_std=None, epsilon=0.05): + return len(self.sec_df_loc) + len(self.sec_cf_loc) + + def fit( + self, + obs_loc: np.ndarray, + obs_B: np.ndarray, + obs_std: np.ndarray | None = None, + epsilon: float = 0.05, + ) -> "SECS": """Fits the SECS to the given observations. Given a number of observation locations and measurements, @@ -81,18 +85,18 @@ def fit(self, obs_loc, obs_B, obs_std=None, epsilon=0.05): Parameters ---------- - obs_locs : ndarray (nobs x 3 [lat, lon, r]) + obs_locs : ndarray (nobs, 3 [lat, lon, r]) Contains latitude, longitude, and radius of the observation locations (place where the measurements are made) - obs_B: ndarray (ntimes x nobs x 3 [Bx, By, Bz]) + obs_B: ndarray (ntimes, nobs, 3 [Bx, By, Bz]) An array containing the measured/observed B-fields. - obs_std : ndarray (ntimes x nobs x 3 [varX, varY, varZ]), optional + obs_std : ndarray (ntimes, nobs, 3 [varX, varY, varZ]), optional Standard error of vector components at each observation location. This can be used to weight different observations more/less heavily. An infinite value eliminates the observation from the fit. - Default: ones(nobs x 3) equal weights + Default: ones(nobs, 3) equal weights epsilon : float Value used to regularize/smooth the SECS amplitudes. Multiplied by the @@ -121,18 +125,17 @@ def fit(self, obs_loc, obs_B, obs_std=None, epsilon=0.05): # Calculate the singular value decomposition (SVD) # NOTE: T_obs has shape (nobs, 3, nsec), we reshape it - # to (nobs*3, nsec); obs_std has shape (ntimes, nobs, 3), + # to (nobs*3, nsec); obs_std has shape (ntimes, nobs, 3), # we reshape it to (ntimes, nobs*3), then loop over ntimes - # to solve using (potentially) time-dependent observation + # to solve using (potentially) time-dependent observation # standard errors to weight the observations for i in range(ntimes): - # Only (re-)calculate SVD when necessary - if i == 0 or not np.all(obs_std[i] == obs_std[i-1]): - + if i == 0 or not np.all(obs_std[i] == obs_std[i - 1]): # Weight T_obs with obs_std - svd_in = (T_obs.reshape(-1, self.nsec) / - obs_std[i].ravel()[:, np.newaxis]) + svd_in = ( + T_obs.reshape(-1, self.nsec) / obs_std[i].ravel()[:, np.newaxis] + ) # Find singular value decompostion U, S, Vh = np.linalg.svd(svd_in, full_matrices=False) @@ -141,31 +144,31 @@ def fit(self, obs_loc, obs_B, obs_std=None, epsilon=0.05): # reciprocal to zero (setting S to infinity firsts avoids # divide-by-zero warings) S[S < epsilon * S.max()] = np.inf - W = 1./S + W = 1.0 / S # Update VWU if obs_std changed VWU = Vh.T @ (np.diag(W) @ U.T) # Solve for SEC amplitudes and error variances - # shape: ntimes x nsec + # shape: (ntimes, nsec) self.sec_amps[i, :] = (VWU @ (obs_B[i] / obs_std[i]).reshape(-1).T).T # Maybe we want the variance of the predictions sometime later...? - # shape: ntimes x nsec + # shape: (ntimes, nsec) valid = np.isfinite(obs_std[i].reshape(-1)) self.sec_amps_var[i, :] = np.sum( - (VWU[:,valid] * obs_std[i].reshape(-1)[valid])**2, - axis=1) + (VWU[:, valid] * obs_std[i].reshape(-1)[valid]) ** 2, axis=1 + ) return self - def fit_unit_currents(self): - """Sets all SECs to a unit current amplitude.""" + def fit_unit_currents(self) -> "SECS": + """Set all SECs to a unit current amplitude.""" self.sec_amps = np.ones((1, self.nsec)) return self - def predict(self, pred_loc, J=False): + def predict(self, pred_loc: np.ndarray, J: bool = False) -> np.ndarray: """Calculate the predicted magnetic field or currents. After a set of observations has been fit to this system we can @@ -174,7 +177,7 @@ def predict(self, pred_loc, J=False): Parameters ---------- - pred_loc: ndarray (npred x 3 [lat, lon, r]) + pred_loc: ndarray (npred, 3 [lat, lon, r]) An array containing the locations where the predictions are desired. J: boolean @@ -183,19 +186,21 @@ def predict(self, pred_loc, J=False): Returns ------- - ndarray (ntimes x npred x 3 [lat, lon, r]) + ndarray (ntimes, npred, 3 [lat, lon, r]) The predicted values calculated from the current amplitudes that were fit to this system. """ if pred_loc.shape[-1] != 3: raise ValueError("Prediction locations must have 3 columns (lat, lon, r)") - if self.sec_amps is None: - raise ValueError("There are no currents associated with the SECs," + - "you need to call .fit() first to fit to some observations.") + if len(self.sec_amps) == 0: + raise ValueError( + "There are no currents associated with the SECs, you need" + "to call .fit() first to fit to some observations." + ) - # T_pred shape=(npred x 3 x nsec) - # sec_amps shape=(nsec x ntimes) + # T_pred shape: (npred, 3, nsec) + # sec_amps shape: (ntimes, nsec) if J: # Predicting currents T_pred = self._calc_J(pred_loc) @@ -207,11 +212,11 @@ def predict(self, pred_loc, J=False): # Therefore this is implemented as tensordot, and the arguments are # arranged to eliminate needs of transposing things later. # The dot product is done over the SEC locations, so the final output - # is of shape: (ntimes x npred x 3) + # is of shape: (ntimes, npred, 3) return np.squeeze(np.tensordot(self.sec_amps, T_pred, (1, 2))) - def predict_B(self, pred_loc): + def predict_B(self, pred_loc: np.ndarray) -> np.ndarray: """Calculate the predicted magnetic fields. After a set of observations has been fit to this system we can @@ -220,18 +225,18 @@ def predict_B(self, pred_loc): Parameters ---------- - pred_loc: ndarray (npred x 3 [lat, lon, r]) + pred_loc: ndarray (npred, 3 [lat, lon, r]) An array containing the locations where the predictions are desired. Returns ------- - ndarray (ntimes x npred x 3 [lat, lon, r]) + ndarray (ntimes, npred, 3 [lat, lon, r]) The predicted values calculated from the current amplitudes that were fit to this system. """ return self.predict(pred_loc) - def predict_J(self, pred_loc): + def predict_J(self, pred_loc: np.ndarray) -> np.ndarray: """Calculate the predicted currents. After a set of observations has been fit to this system we can @@ -240,18 +245,18 @@ def predict_J(self, pred_loc): Parameters ---------- - pred_loc: ndarray (npred x 3 [lat, lon, r]) + pred_loc: ndarray (npred, 3 [lat, lon, r]) An array containing the locations where the predictions are desired. Returns ------- - ndarray (ntimes x npred x 3 [lat, lon, r]) + ndarray (ntimes, npred, 3 [lat, lon, r]) The predicted values calculated from the current amplitudes that were fit to this system. """ return self.predict(pred_loc, J=True) - def _calc_T(self, obs_loc): + def _calc_T(self, obs_loc: np.ndarray) -> np.ndarray: """Calculate the T transfer matrix. The magnetic field transfer matrix to go from SEC locations to observation @@ -271,7 +276,7 @@ def _calc_T(self, obs_loc): return T - def _calc_J(self, obs_loc): + def _calc_J(self, obs_loc: np.ndarray) -> np.ndarray: """Calculate the J transfer matrix. The current transfer matrix to go from SEC locations to observation @@ -292,7 +297,7 @@ def _calc_J(self, obs_loc): return J -def T_df(obs_loc, sec_loc): +def T_df(obs_loc: np.ndarray, sec_loc: np.ndarray) -> np.ndarray: """Calculate the divergence free magnetic field transfer function. The transfer function goes from SEC location to observation location @@ -321,25 +326,26 @@ def T_df(obs_loc, sec_loc): alpha = calc_bearing(obs_loc[:, :2], sec_loc[:, :2]) # magnetic permeability - mu0 = 4*np.pi*1e-7 + mu0 = 4 * np.pi * 1e-7 # simplify calculations by storing this ratio - x = obs_r/sec_r + x = obs_r / sec_r sin_theta = np.sin(theta) cos_theta = np.cos(theta) - factor = 1./np.sqrt(1 - 2*x*cos_theta + x**2) + factor = 1.0 / np.sqrt(1 - 2 * x * cos_theta + x**2) # Amm & Viljanen: Equation 9 - Br = mu0/(4*np.pi*obs_r) * (factor - 1) + Br = mu0 / (4 * np.pi * obs_r) * (factor - 1) # Amm & Viljanen: Equation 10 (transformed to try and eliminate trig operations and # divide by zeros) - Btheta = -mu0/(4*np.pi*obs_r) * (factor*(x - cos_theta) + cos_theta) + Btheta = -mu0 / (4 * np.pi * obs_r) * (factor * (x - cos_theta) + cos_theta) # If sin(theta) == 0: Btheta = 0 # There is a possible 0/0 in the expansion when sec_loc == obs_loc - Btheta = np.divide(Btheta, sin_theta, out=np.zeros_like(sin_theta), - where=sin_theta != 0) + Btheta = np.divide( + Btheta, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0 + ) # When observation points radii are outside of the sec locations under_locs = sec_r < obs_r @@ -350,18 +356,29 @@ def T_df(obs_loc, sec_loc): # except for very large matrices. if np.any(under_locs): # Flipped from previous case - x = sec_r/obs_r + x = sec_r / obs_r # Amm & Viljanen: Equation A.7 - Br2 = mu0*x/(4*np.pi*obs_r) * (1./np.sqrt(1 - 2*x*cos_theta + x**2) - 1) + Br2 = ( + mu0 + * x + / (4 * np.pi * obs_r) + * (1.0 / np.sqrt(1 - 2 * x * cos_theta + x**2) - 1) + ) # Amm & Viljanen: Equation A.8 - Btheta2 = - mu0 / (4*np.pi*obs_r) * ((obs_r-sec_r*cos_theta) / - np.sqrt(obs_r**2 - - 2*obs_r*sec_r*cos_theta + - sec_r**2) - 1) - Btheta2 = np.divide(Btheta2, sin_theta, out=np.zeros_like(sin_theta), - where=sin_theta != 0) + Btheta2 = ( + -mu0 + / (4 * np.pi * obs_r) + * ( + (obs_r - sec_r * cos_theta) + / np.sqrt(obs_r**2 - 2 * obs_r * sec_r * cos_theta + sec_r**2) + - 1 + ) + ) + Btheta2 = np.divide( + Btheta2, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0 + ) # Update only the locations where secs are under observations Btheta[under_locs] = Btheta2[under_locs] @@ -370,14 +387,14 @@ def T_df(obs_loc, sec_loc): # Transform back to Bx, By, Bz at each local point T = np.empty((nobs, 3, nsec)) # alpha == angle (from cartesian x-axis (By), going towards y-axis (Bx)) - T[:, 0, :] = -Btheta*np.sin(alpha) - T[:, 1, :] = -Btheta*np.cos(alpha) + T[:, 0, :] = -Btheta * np.sin(alpha) + T[:, 1, :] = -Btheta * np.cos(alpha) T[:, 2, :] = -Br return T -def T_cf(obs_loc, sec_loc): +def T_cf(obs_loc: np.ndarray, sec_loc: np.ndarray) -> np.ndarray: """Calculate the curl free magnetic field transfer function. The transfer function goes from SEC location to observation location @@ -396,10 +413,12 @@ def T_cf(obs_loc, sec_loc): ndarray (nobs, 3, nsec) The T transfer matrix. """ - raise NotImplementedError("Curl Free Magnetic Field Transfers are not implemented yet.") + raise NotImplementedError( + "Curl Free Magnetic Field Transfers " "are not implemented yet." + ) -def J_df(obs_loc, sec_loc): +def J_df(obs_loc: np.ndarray, sec_loc: np.ndarray) -> np.ndarray: """Calculate the divergence free current density transfer function. The transfer function goes from SEC location to observation location @@ -429,25 +448,29 @@ def J_df(obs_loc, sec_loc): alpha = calc_bearing(obs_loc[:, :2], sec_loc[:, :2]) # Amm & Viljanen: Equation 6 - tan_theta2 = np.tan(theta/2.) - - J_phi = 1./(4*np.pi*sec_r) - J_phi = np.divide(J_phi, tan_theta2, out=np.ones_like(tan_theta2)*np.inf, - where=tan_theta2 != 0.) + tan_theta2 = np.tan(theta / 2.0) + + J_phi = 1.0 / (4 * np.pi * sec_r) + J_phi = np.divide( + J_phi, + tan_theta2, + out=np.ones_like(tan_theta2) * np.inf, + where=tan_theta2 != 0.0, + ) # Only valid on the SEC shell - J_phi[sec_r != obs_r] = 0. + J_phi[sec_r != obs_r] = 0.0 # Transform back to Bx, By, Bz at each local point J = np.empty((nobs, 3, nsec)) # alpha == angle (from cartesian x-axis (By), going towards y-axis (Bx)) - J[:, 0, :] = -J_phi*np.cos(alpha) - J[:, 1, :] = J_phi*np.sin(alpha) - J[:, 2, :] = 0. + J[:, 0, :] = -J_phi * np.cos(alpha) + J[:, 1, :] = J_phi * np.sin(alpha) + J[:, 2, :] = 0.0 return J -def J_cf(obs_loc, sec_loc): +def J_cf(obs_loc: np.ndarray, sec_loc: np.ndarray) -> np.ndarray: """Calculate the curl free magnetic field transfer function. The transfer function goes from SEC location to observation location @@ -476,32 +499,36 @@ def J_cf(obs_loc, sec_loc): alpha = calc_bearing(obs_loc[:, :2], sec_loc[:, :2]) # Amm & Viljanen: Equation 7 - tan_theta2 = np.tan(theta/2.) - - J_theta = 1./(4*np.pi*sec_r) - J_theta = np.divide(J_theta, tan_theta2, out=np.ones_like(tan_theta2)*np.inf, - where=tan_theta2 != 0) + tan_theta2 = np.tan(theta / 2.0) + + J_theta = 1.0 / (4 * np.pi * sec_r) + J_theta = np.divide( + J_theta, + tan_theta2, + out=np.ones_like(tan_theta2) * np.inf, + where=tan_theta2 != 0, + ) # Uniformly directed FACs around the globe, except the pole # Integrated over the globe, this will lead to zero - J_r = -np.ones(J_theta.shape)/(4*np.pi*sec_r**2) - J_r[theta == 0.] = 1. + J_r = -np.ones(J_theta.shape) / (4 * np.pi * sec_r**2) + J_r[theta == 0.0] = 1.0 # Only valid on the SEC shell - J_theta[sec_r != obs_r] = 0. - J_r[sec_r != obs_r] = 0. + J_theta[sec_r != obs_r] = 0.0 + J_r[sec_r != obs_r] = 0.0 # Transform back to Bx, By, Bz at each local point J = np.empty((nobs, 3, nsec)) # alpha == angle (from cartesian x-axis (By), going towards y-axis (Bx)) - J[:, 0, :] = -J_theta*np.sin(alpha) - J[:, 1, :] = -J_theta*np.cos(alpha) + J[:, 0, :] = -J_theta * np.sin(alpha) + J[:, 1, :] = -J_theta * np.cos(alpha) J[:, 2, :] = -J_r return J -def calc_angular_distance(latlon1, latlon2): +def calc_angular_distance(latlon1: np.ndarray, latlon2: np.ndarray) -> np.ndarray: """Calculate the angular distance between a set of points. This function calculates the angular distance in radians @@ -509,15 +536,15 @@ def calc_angular_distance(latlon1, latlon2): Parameters ---------- - latlon1 : ndarray (n x 2 [lat, lon]) + latlon1 : ndarray (n, 2 [lat, lon]) An array of n (latitude, longitude) points. - latlon2 : ndarray (m x 2 [lat, lon]) + latlon2 : ndarray (m, 2 [lat, lon]) An array of m (latitude, longitude) points. Returns ------- - ndarray (n x m) + ndarray (n, m) The array of distances between the input arrays. """ lat1 = np.deg2rad(latlon1[:, 0])[:, np.newaxis] @@ -528,12 +555,13 @@ def calc_angular_distance(latlon1, latlon2): dlon = lon2 - lon1 # theta == angular distance between two points - theta = np.arccos(np.sin(lat1)*np.sin(lat2) + - np.cos(lat1)*np.cos(lat2)*np.cos(dlon)) + theta = np.arccos( + np.sin(lat1) * np.sin(lat2) + np.cos(lat1) * np.cos(lat2) * np.cos(dlon) + ) return theta -def calc_bearing(latlon1, latlon2): +def calc_bearing(latlon1: np.ndarray, latlon2: np.ndarray) -> np.ndarray: """Calculate the bearing (direction) between a set of points. This function calculates the bearing in radians @@ -543,15 +571,15 @@ def calc_bearing(latlon1, latlon2): Parameters ---------- - latlon1 : ndarray (n x 2 [lat, lon]) + latlon1 : ndarray (n, 2 [lat, lon]) An array of n (latitude, longitude) points. - latlon2 : ndarray (m x 2 [lat, lon]) + latlon2 : ndarray (m, 2 [lat, lon]) An array of m (latitude, longitude) points. Returns ------- - ndarray (n x m) + ndarray (n, m) The array of bearings between the input arrays. """ lat1 = np.deg2rad(latlon1[:, 0])[:, np.newaxis] @@ -568,7 +596,8 @@ def calc_bearing(latlon1, latlon2): # SEC coordinates are: theta (colatitude (+ away from North Pole)), # phi (longitude, + east), r (+ out) # Obs coordinates are: X (+ north), Y (+ east), Z (+ down) - alpha = np.pi/2 - np.arctan2(np.sin(dlon)*np.cos(lat2), - np.cos(lat1)*np.sin(lat2) - - np.sin(lat1)*np.cos(lat2)*np.cos(dlon)) + alpha = np.pi / 2 - np.arctan2( + np.sin(dlon) * np.cos(lat2), + np.cos(lat1) * np.sin(lat2) - np.sin(lat1) * np.cos(lat2) * np.cos(dlon), + ) return alpha diff --git a/tests/test_secs.py b/tests/test_secs.py index 8a3b973..e13b72e 100644 --- a/tests/test_secs.py +++ b/tests/test_secs.py @@ -1,41 +1,55 @@ import numpy as np -from numpy.testing import assert_array_equal, assert_allclose import pytest +from numpy.testing import assert_allclose, assert_array_equal + import pysecs + R_EARTH = 6378e3 def test_angular_distance(): "Test the angular distance formula." - latlon1 = np.array([[0., 0.]]) - latlon2 = np.array([[0., 0.], [0., 90], [-90., 0.], [0., 180.]]) - assert_array_equal(pysecs.calc_angular_distance(latlon1, latlon2), - np.deg2rad([[0., 90., 90., 180.]])) + latlon1 = np.array([[0.0, 0.0]]) + latlon2 = np.array([[0.0, 0.0], [0.0, 90], [-90.0, 0.0], [0.0, 180.0]]) + assert_array_equal( + pysecs.calc_angular_distance(latlon1, latlon2), + np.deg2rad([[0.0, 90.0, 90.0, 180.0]]), + ) def test_bearing(): "Test the cardinal directions." - latlon1 = np.array([[0., 0.]]) - latlon2 = np.array([[0., 90.], [90., 0.], [90., 45.], [0., -90.], [-90, 0.]]) - assert_array_equal(pysecs.calc_bearing(latlon1, latlon2), - np.deg2rad([[0., 90., 90., 180., -90]])) + latlon1 = np.array([[0.0, 0.0]]) + latlon2 = np.array( + [[0.0, 90.0], [90.0, 0.0], [90.0, 45.0], [0.0, -90.0], [-90, 0.0]] + ) + assert_array_equal( + pysecs.calc_bearing(latlon1, latlon2), + np.deg2rad([[0.0, 90.0, 90.0, 180.0, -90]]), + ) def test_divergence_free_magnetic_directions(): "Make sure the divergence free magnetic field angles are correct" # Place the SEC at the equator sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going around in a circle from the point - obs_latlonr = np.array([[5., 0., R_EARTH], [0., 5., R_EARTH], - [-5, 0., R_EARTH], [0., -5., R_EARTH]]) + obs_latlonr = np.array( + [ + [5.0, 0.0, R_EARTH], + [0.0, 5.0, R_EARTH], + [-5, 0.0, R_EARTH], + [0.0, -5.0, R_EARTH], + ] + ) B = np.squeeze(pysecs.T_df(obs_latlonr, sec_latlonr)) angles = np.arctan2(B[:, 0], B[:, 1]) # southward, westward, northward, eastward - expected_angles = np.deg2rad([-90, 180., 90., 0.]) + expected_angles = np.deg2rad([-90, 180.0, 90.0, 0.0]) assert_allclose(angles, expected_angles, rtol=1e-10, atol=1e-10) @@ -43,11 +57,11 @@ def test_divergence_free_magnetic_magnitudes_obs_under(): "Make sure the divergence free magnetic amplitudes are correct." # Place the SEC at the North Pole sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going out in an angle from the SEC (in longitude) angles = np.linspace(0.1, 180) obs_r = R_EARTH - obs_latlonr = np.zeros(angles.shape + (3,)) + obs_latlonr = np.zeros((*angles.shape, 3)) obs_latlonr[:, 1] = angles obs_latlonr[:, 2] = obs_r @@ -58,23 +72,25 @@ def test_divergence_free_magnetic_magnitudes_obs_under(): assert_allclose(np.zeros(angles.shape), B[:, 0], atol=1e-16) # Actual magnitude - mu0 = 4*np.pi*1e-7 + mu0 = 4 * np.pi * 1e-7 # simplify calculations by storing this ratio - x = obs_r/sec_r + x = obs_r / sec_r sin_theta = np.sin(np.deg2rad(angles)) cos_theta = np.cos(np.deg2rad(angles)) - factor = 1./np.sqrt(1 - 2*x*cos_theta + x**2) + factor = 1.0 / np.sqrt(1 - 2 * x * cos_theta + x**2) # Amm & Viljanen: Equation 9 - Br = mu0/(4*np.pi*obs_r) * (factor - 1) + Br = mu0 / (4 * np.pi * obs_r) * (factor - 1) # Bz in opposite direction of Br assert_allclose(-Br, B[:, 2]) # Amm & Viljanen: Equation 10 - Btheta = -mu0/(4*np.pi*obs_r) * (factor*(x - cos_theta) + cos_theta) - Btheta = np.divide(Btheta, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0) + Btheta = -mu0 / (4 * np.pi * obs_r) * (factor * (x - cos_theta) + cos_theta) + Btheta = np.divide( + Btheta, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0 + ) assert_allclose(Btheta, B[:, 1]) @@ -82,11 +98,11 @@ def test_divergence_free_magnetic_magnitudes_obs_over(): "Make sure the divergence free magnetic amplitudes are correct." # Place the SEC at the North Pole sec_r = R_EARTH - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going out in an angle from the SEC (in longitude) angles = np.linspace(0.1, 180) obs_r = R_EARTH + 100 - obs_latlonr = np.zeros(angles.shape + (3,)) + obs_latlonr = np.zeros((*angles.shape, 3)) obs_latlonr[:, 1] = angles obs_latlonr[:, 2] = obs_r @@ -97,54 +113,75 @@ def test_divergence_free_magnetic_magnitudes_obs_over(): assert_allclose(np.zeros(angles.shape), B[:, 0], atol=1e-16) # Actual magnitude - mu0 = 4*np.pi*1e-7 - x = sec_r/obs_r + mu0 = 4 * np.pi * 1e-7 + x = sec_r / obs_r sin_theta = np.sin(np.deg2rad(angles)) cos_theta = np.cos(np.deg2rad(angles)) # Amm & Viljanen: Equation A.7 - Br = mu0*x/(4*np.pi*obs_r) * (1./np.sqrt(1 - 2*x*cos_theta + x**2) - 1) + Br = ( + mu0 + * x + / (4 * np.pi * obs_r) + * (1.0 / np.sqrt(1 - 2 * x * cos_theta + x**2) - 1) + ) # Bz in opposite direction of Br assert_allclose(-Br, B[:, 2]) # Amm & Viljanen: Equation A.8 - Btheta = -mu0/(4*np.pi*obs_r)*((obs_r-sec_r*cos_theta) / - np.sqrt(obs_r**2 - 2*obs_r*sec_r*cos_theta + sec_r**2) - 1) - Btheta = np.divide(Btheta, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0) + Btheta = ( + -mu0 + / (4 * np.pi * obs_r) + * ( + (obs_r - sec_r * cos_theta) + / np.sqrt(obs_r**2 - 2 * obs_r * sec_r * cos_theta + sec_r**2) + - 1 + ) + ) + Btheta = np.divide( + Btheta, sin_theta, out=np.zeros_like(sin_theta), where=sin_theta != 0 + ) assert_allclose(Btheta, B[:, 1]) def test_outside_current_plane(): "Make sure all currents outside the SEC plane are 0." sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Above and below the plane, also on and off the SEC point - obs_latlonr = np.array([[0., 0., sec_r - 100.], [0., 0., sec_r + 100.], - [5, 0., sec_r - 100.], [5., 0., sec_r + 100.]]) + obs_latlonr = np.array( + [ + [0.0, 0.0, sec_r - 100.0], + [0.0, 0.0, sec_r + 100.0], + [5, 0.0, sec_r - 100.0], + [5.0, 0.0, sec_r + 100.0], + ] + ) # df currents J = np.squeeze(pysecs.J_df(obs_latlonr, sec_latlonr)) - assert np.all(J == 0.) + assert np.all(J == 0.0) # cf currents J = np.squeeze(pysecs.J_cf(obs_latlonr, sec_latlonr)) - assert np.all(J == 0.) + assert np.all(J == 0.0) def test_divergence_free_current_directions(): "Make sure the divergence free current angles are correct." # Place the SEC at the equator sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going around in a circle from the point - obs_latlonr = np.array([[5., 0., sec_r], [0., 5., sec_r], - [-5, 0., sec_r], [0., -5., sec_r]]) + obs_latlonr = np.array( + [[5.0, 0.0, sec_r], [0.0, 5.0, sec_r], [-5, 0.0, sec_r], [0.0, -5.0, sec_r]] + ) J = np.squeeze(pysecs.J_df(obs_latlonr, sec_latlonr)) angles = np.arctan2(J[:, 0], J[:, 1]) # westward, northward, eastward, southward - expected_angles = np.deg2rad([-180., 90., 0., -90.]) + expected_angles = np.deg2rad([-180.0, 90.0, 0.0, -90.0]) assert_allclose(angles, expected_angles, atol=1e-16) @@ -152,27 +189,31 @@ def test_divergence_free_current_magnitudes(): "Make sure the divergence free current amplitudes are correct." # Place the SEC at the North Pole sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going out in an angle from the SEC (in longitude) angles = np.linspace(0.1, 180) - obs_latlonr = np.zeros(angles.shape + (3,)) + obs_latlonr = np.zeros((*angles.shape, 3)) obs_latlonr[:, 1] = angles obs_latlonr[:, 2] = sec_r J = np.squeeze(pysecs.J_df(obs_latlonr, sec_latlonr)) # Make sure all radial components are zero in this system - assert np.all(J[:, 2] == 0.) + assert np.all(J[:, 2] == 0.0) # Also all y components (angles goes around the equator and all # quantities should be perpendicular to that) assert_allclose(np.zeros(angles.shape), J[:, 1], atol=1e-16) # Actual magnitude - tan_theta2 = np.tan(np.deg2rad(angles/2)) - J_test = 1./(4*np.pi*sec_r) - J_test = np.divide(J_test, tan_theta2, out=np.ones_like(tan_theta2)*np.inf, - where=tan_theta2 != 0.) + tan_theta2 = np.tan(np.deg2rad(angles / 2)) + J_test = 1.0 / (4 * np.pi * sec_r) + J_test = np.divide( + J_test, + tan_theta2, + out=np.ones_like(tan_theta2) * np.inf, + where=tan_theta2 != 0.0, + ) assert_allclose(J_test, J[:, 0], atol=1e-16) @@ -181,17 +222,18 @@ def test_curl_free_current_directions(): "Make sure the curl free current angles are correct." # Place the SEC at the equator sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going around in a circle from the point - obs_latlonr = np.array([[5., 0., sec_r], [0., 5., sec_r], - [-5, 0., sec_r], [0., -5., sec_r]]) + obs_latlonr = np.array( + [[5.0, 0.0, sec_r], [0.0, 5.0, sec_r], [-5, 0.0, sec_r], [0.0, -5.0, sec_r]] + ) J = np.squeeze(pysecs.J_cf(obs_latlonr, sec_latlonr)) angles = np.arctan2(J[:, 0], J[:, 1]) # pointing out from the SEC direction to OBS direction. # northward, eastward, southward, westward - expected_angles = np.deg2rad([90., 0., -90., -180]) + expected_angles = np.deg2rad([90.0, 0.0, -90.0, -180]) assert_allclose(angles, expected_angles, atol=1e-15) @@ -199,17 +241,17 @@ def test_curl_free_current_magnitudes(): "Make sure the curl free current amplitudes are correct." # Place the SEC at the North Pole sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going out in an angle from the SEC (in longitude) angles = np.linspace(0.1, 180) - obs_latlonr = np.zeros(angles.shape + (3,)) + obs_latlonr = np.zeros((*angles.shape, 3)) obs_latlonr[:, 1] = angles obs_latlonr[:, 2] = sec_r J = np.squeeze(pysecs.J_cf(obs_latlonr, sec_latlonr)) # Make sure all radial components are oppositely directed - radial_component = 1./(4*np.pi*sec_r**2) + radial_component = 1.0 / (4 * np.pi * sec_r**2) assert np.all(J[:, 2] == radial_component) # All x components should be zero (angles goes around the equator and all @@ -218,10 +260,14 @@ def test_curl_free_current_magnitudes(): assert_allclose(np.zeros(angles.shape), J[:, 0], atol=1e-16) # Actual magnitude - tan_theta2 = np.tan(np.deg2rad(angles/2)) - J_test = 1./(4*np.pi*sec_r) - J_test = np.divide(J_test, tan_theta2, out=np.ones_like(tan_theta2)*np.inf, - where=tan_theta2 != 0.) + tan_theta2 = np.tan(np.deg2rad(angles / 2)) + J_test = 1.0 / (4 * np.pi * sec_r) + J_test = np.divide( + J_test, + tan_theta2, + out=np.ones_like(tan_theta2) * np.inf, + where=tan_theta2 != 0.0, + ) assert_allclose(J_test, J[:, 1], atol=1e-16) @@ -231,26 +277,26 @@ def test_curl_free_magnetic_magnitudes(): # TODO: Update this test once T_cf is implemented # Place the SEC at the North Pole sec_r = R_EARTH + 100 - sec_latlonr = np.array([[0., 0., sec_r]]) + sec_latlonr = np.array([[0.0, 0.0, sec_r]]) # Going out in an angle from the SEC (in longitude) angles = np.linspace(0.1, 180) - obs_latlonr = np.zeros(angles.shape + (3,)) + obs_latlonr = np.zeros((*angles.shape, 3)) obs_latlonr[:, 1] = angles obs_latlonr[:, 2] = sec_r - with pytest.raises(NotImplementedError, match='Curl Free Magnetic'): + with pytest.raises(NotImplementedError, match="Curl Free Magnetic"): pysecs.T_cf(obs_latlonr, sec_latlonr) def test_empty_object(): "Testing empty secs object creation failure." - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Must specify"): pysecs.SECS() def test_list_numpy(): "Make sure creation with numpy and list produce the same locations." - x2d = [[1., 0., 0.], [1., 0., 0.]] + x2d = [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]] x2d_np = np.array(x2d) secs_list2 = pysecs.SECS(sec_df_loc=x2d, sec_cf_loc=x2d) secs_np2 = pysecs.SECS(sec_df_loc=x2d_np, sec_cf_loc=x2d_np) @@ -264,9 +310,9 @@ def test_sec_bad_shape(): """Test bad input shape.""" # Wrong dimensions x = np.array([[1, 0], [1, 0]]) - with pytest.raises(ValueError, match='SEC DF locations'): + with pytest.raises(ValueError, match="SEC DF locations"): pysecs.SECS(sec_df_loc=x) - with pytest.raises(ValueError, match='SEC CF locations'): + with pytest.raises(ValueError, match="SEC CF locations"): pysecs.SECS(sec_cf_loc=x) @@ -282,25 +328,27 @@ def test_one_sec(): def test_fit_unit_currents(): """Test the unit current function.""" # divergence free - secs = pysecs.SECS(sec_df_loc=[[1., 0., 0.], [1., 0., 0.]]) + secs = pysecs.SECS(sec_df_loc=[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) secs.fit_unit_currents() assert_array_equal(np.ones((1, 2)), secs.sec_amps) # curl free - secs = pysecs.SECS(sec_cf_loc=[[1., 0., 0.], [1., 0., 0.]]) + secs = pysecs.SECS(sec_cf_loc=[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) secs.fit_unit_currents() assert_array_equal(np.ones((1, 2)), secs.sec_amps) # divergence free + curl free - secs = pysecs.SECS(sec_df_loc=[[1., 0., 0.], [1., 0., 0.]], - sec_cf_loc=[[1., 0., 0.], [1., 0., 0.]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + sec_cf_loc=[[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + ) secs.fit_unit_currents() assert_array_equal(np.ones((1, 4)), secs.sec_amps) def test_fit_bad_obsloc(): """Test bad observation locations input.""" - secs = pysecs.SECS(sec_df_loc=[1., 0., R_EARTH + 1e6]) + secs = pysecs.SECS(sec_df_loc=[1.0, 0.0, R_EARTH + 1e6]) obs_loc = np.array([[0, 0]]) obs_B = np.array([[1, 1, 1]]) with pytest.raises(ValueError, match="Observation locations"): @@ -309,31 +357,34 @@ def test_fit_bad_obsloc(): def test_fit_one_time(): """One timestep fit.""" - secs = pysecs.SECS(sec_df_loc=[[1., 0., R_EARTH + 1e6], - [-1., 0., R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, R_EARTH + 1e6], [-1.0, 0.0, R_EARTH + 1e6]] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.array([[1, 1, 1]]) secs.fit(obs_loc, obs_B) - assert_allclose([[6.40594202e+13, -7.41421248e+13]], secs.sec_amps) + assert_allclose([[6.40594202e13, -7.41421248e13]], secs.sec_amps) def test_fit_multi_time(): """Multi timestep fits.""" - secs = pysecs.SECS(sec_df_loc=[[1., 0., R_EARTH + 1e6], - [-1., 0., R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, R_EARTH + 1e6], [-1.0, 0.0, R_EARTH + 1e6]] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.ones((2, 1, 3)) obs_B[1, :, :] *= 2 secs.fit(obs_loc, obs_B) - arr = np.array([6.40594202e+13, -7.41421248e+13]) - expected = np.array([arr, 2*arr]) + arr = np.array([6.40594202e13, -7.41421248e13]) + expected = np.array([arr, 2 * arr]) assert_allclose(expected, secs.sec_amps) def test_fit_obs_std(): """Test that variance on observations changes the results.""" - secs = pysecs.SECS(sec_df_loc=[[1., 0., R_EARTH + 1e6], - [-1., 0., R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, R_EARTH + 1e6], [-1.0, 0.0, R_EARTH + 1e6]] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.ones((2, 1, 3)) obs_B[1, :, :] *= 2 @@ -341,15 +392,15 @@ def test_fit_obs_std(): # Remove the z component from the fit of the second timestep obs_std[1, :, 2] = np.inf secs.fit(obs_loc, obs_B, obs_std=obs_std) - expected = np.array([[6.40594202e+13, -7.41421248e+13], - [1.382015e+14, -1.382015e+14]]) + expected = np.array([[6.40594202e13, -7.41421248e13], [1.382015e14, -1.382015e14]]) assert_allclose(expected, secs.sec_amps, rtol=1e-6) def test_fit_epsilon(): """Test that epsilon removes some components.""" - secs = pysecs.SECS(sec_df_loc=[[1., 0., R_EARTH + 1e6], - [-1., 0., R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, R_EARTH + 1e6], [-1.0, 0.0, R_EARTH + 1e6]] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.ones((2, 1, 3)) obs_B[1, :, :] *= 2 @@ -357,15 +408,15 @@ def test_fit_epsilon(): # Remove the z component from the fit of the second timestep obs_std[1, :, 2] = np.inf secs.fit(obs_loc, obs_B, obs_std=obs_std, epsilon=0.8) - expected = np.array([[-5.041352e+12, -5.041352e+12], - [1.382015e+14, -1.382015e+14]]) + expected = np.array([[-5.041352e12, -5.041352e12], [1.382015e14, -1.382015e14]]) assert_allclose(expected, secs.sec_amps, rtol=1e-6) def test_bad_predict(): """Test bad input to predict.""" - secs = pysecs.SECS(sec_df_loc=[[1., 0., R_EARTH + 1e6], - [-1., 0., R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[[1.0, 0.0, R_EARTH + 1e6], [-1.0, 0.0, R_EARTH + 1e6]] + ) # Calling predict with the wrong shape pred_loc = np.array([[0, 0]]) @@ -380,14 +431,18 @@ def test_bad_predict(): def test_predictB(): """Test that epsilon removes some components.""" - secs = pysecs.SECS(sec_df_loc=[[1, 0, R_EARTH + 1e6], - [-1, 0, R_EARTH + 1e6], - [-1, 1, R_EARTH + 1e6], - [1, 1, R_EARTH + 1e6], - [0, 1, R_EARTH + 1e6], - [0, -1, R_EARTH + 1e6], - [-1, -1, R_EARTH + 1e6], - [1, -1, R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[ + [1, 0, R_EARTH + 1e6], + [-1, 0, R_EARTH + 1e6], + [-1, 1, R_EARTH + 1e6], + [1, 1, R_EARTH + 1e6], + [0, 1, R_EARTH + 1e6], + [0, -1, R_EARTH + 1e6], + [-1, -1, R_EARTH + 1e6], + [1, -1, R_EARTH + 1e6], + ] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.ones((2, 1, 3)) obs_B[1, :, :] *= 2 @@ -403,14 +458,18 @@ def test_predictB(): def test_predictJ(): """Test the current sheet predictions (df).""" - secs = pysecs.SECS(sec_df_loc=[[1, 0, R_EARTH + 1e6], - [-1, 0, R_EARTH + 1e6], - [-1, 1, R_EARTH + 1e6], - [1, 1, R_EARTH + 1e6], - [0, 1, R_EARTH + 1e6], - [0, -1, R_EARTH + 1e6], - [-1, -1, R_EARTH + 1e6], - [1, -1, R_EARTH + 1e6]]) + secs = pysecs.SECS( + sec_df_loc=[ + [1, 0, R_EARTH + 1e6], + [-1, 0, R_EARTH + 1e6], + [-1, 1, R_EARTH + 1e6], + [1, 1, R_EARTH + 1e6], + [0, 1, R_EARTH + 1e6], + [0, -1, R_EARTH + 1e6], + [-1, -1, R_EARTH + 1e6], + [1, -1, R_EARTH + 1e6], + ] + ) obs_loc = np.array([[0, 0, R_EARTH]]) obs_B = np.ones((2, 1, 3)) obs_B[1, :, :] *= 2 @@ -423,8 +482,8 @@ def test_predictJ(): # Move up to the current sheet pred_loc = np.array([[0, 0, R_EARTH + 1e6]]) J_pred = secs.predict(pred_loc, J=True) - arr = np.array([-1.148475e+08, 1.148417e+08, 0.000000e+00]) - expected = np.array([arr, 2*arr]) + arr = np.array([-1.148475e08, 1.148417e08, 0.000000e00]) + expected = np.array([arr, 2 * arr]) assert_allclose(expected, J_pred, rtol=1e-6) # Use the predict_J function call directly @@ -433,14 +492,18 @@ def test_predictJ(): def test_predictJ_cf(): """Test the current sheet predictions (cf).""" - sec_loc = np.array([[1, 0, R_EARTH + 1e6], - [-1, 0, R_EARTH + 1e6], - [-1, 1, R_EARTH + 1e6], - [1, 1, R_EARTH + 1e6], - [0, 1, R_EARTH + 1e6], - [0, -1, R_EARTH + 1e6], - [-1, -1, R_EARTH + 1e6], - [1, -1, R_EARTH + 1e6]]) + sec_loc = np.array( + [ + [1, 0, R_EARTH + 1e6], + [-1, 0, R_EARTH + 1e6], + [-1, 1, R_EARTH + 1e6], + [1, 1, R_EARTH + 1e6], + [0, 1, R_EARTH + 1e6], + [0, -1, R_EARTH + 1e6], + [-1, -1, R_EARTH + 1e6], + [1, -1, R_EARTH + 1e6], + ] + ) secs = pysecs.SECS(sec_cf_loc=sec_loc) obs_loc = np.array([[0, 0, R_EARTH]]) secs.fit_unit_currents() @@ -452,7 +515,7 @@ def test_predictJ_cf(): # Move up to the current sheet pred_loc = np.array([[0, 0, R_EARTH + 1e6]]) J_pred = secs.predict(pred_loc, J=True) - expected = np.array([0, 0, 1.169507e-14]) + expected = np.array([0, 0, 1.169507e-14]) assert_allclose(expected, J_pred, rtol=1e-6, atol=1e-10) # Use the predict_J function call directly @@ -461,14 +524,18 @@ def test_predictJ_cf(): def test_predictJ_cf_df(): """Test the current sheet predictions (cf+df).""" - sec_loc = np.array([[1, 0, R_EARTH + 1e6], - [-1, 0, R_EARTH + 1e6], - [-1, 1, R_EARTH + 1e6], - [1, 1, R_EARTH + 1e6], - [0, 1, R_EARTH + 1e6], - [0, -1, R_EARTH + 1e6], - [-1, -1, R_EARTH + 1e6], - [1, -1, R_EARTH + 1e6]]) + sec_loc = np.array( + [ + [1, 0, R_EARTH + 1e6], + [-1, 0, R_EARTH + 1e6], + [-1, 1, R_EARTH + 1e6], + [1, 1, R_EARTH + 1e6], + [0, 1, R_EARTH + 1e6], + [0, -1, R_EARTH + 1e6], + [-1, -1, R_EARTH + 1e6], + [1, -1, R_EARTH + 1e6], + ] + ) secs = pysecs.SECS(sec_df_loc=sec_loc, sec_cf_loc=sec_loc) obs_loc = np.array([[0, 0, R_EARTH]]) secs.fit_unit_currents()