import matplotlib.pyplot as plt
from datatree import DataTree, register_datatree_accessor
from numpy import stack
from xarray import Dataset
from swarmpal.io import PalProcess
from swarmpal.toolboxes.fac.fac_algorithms import fac_single_sat_algo
__all__ = (
"FAC_singlesat",
"PalFacDataTreeAccessor",
)
[docs]class FAC_singlesat(PalProcess):
"""Provides the process for the classic single-satellite FAC algorithm
Notes
-----
Expected config parameters:
dataset
model_varname
measurement_varname
"""
@property
def process_name(self):
return "FAC_singlesat"
[docs] def set_config(
self,
dataset: str = "SW_OPER_MAGA_LR_1B",
model_varname: str = "B_NEC_CHAOS",
measurement_varname: str = "B_NEC",
) -> None:
self.config = dict(
dataset=dataset,
model_varname=model_varname,
measurement_varname=measurement_varname,
)
def _call(self, datatree):
# Identify inputs for algorithm
subtree = datatree[self.config.get("dataset")]
dataset = subtree.ds
time = self._get_time(dataset)
positions = self._get_positions(dataset)
B_res = self._get_B_res(dataset)
B_model = self._get_B_model(dataset)
# Apply algorithm
fac_results = fac_single_sat_algo(
time=time, positions=positions, B_res=B_res, B_model=B_model
)
time_out = fac_results["time"]
fac_out = fac_results["fac"]
# Insert a new output dataset with these results
ds_out = Dataset(
data_vars={
"Timestamp": ("Timestamp", time_out),
"FAC": ("Timestamp", fac_out),
}
)
ds_out["FAC"].attrs = {"units": "uA/m2"}
subtree["output"] = DataTree(data=ds_out)
return datatree
def _validate(self):
...
def _get_time(self, dataset):
return dataset.get("Timestamp").data.astype("datetime64[ns]")
def _get_positions(self, dataset):
return stack(
[
dataset.get("Latitude").data,
dataset.get("Longitude").data,
dataset.get("Radius").data,
],
axis=1,
)
def _get_B_res(self, dataset):
measurement_varname = self.config.get("measurement_varname", "B_NEC")
model_varname = self.config.get("model_varname", "B_NEC_Model")
return dataset.get(measurement_varname).data - dataset.get(model_varname).data
def _get_B_model(self, dataset):
model_varname = self.config.get("model_varname", "B_NEC_Model")
return dataset.get(model_varname).data
[docs]@register_datatree_accessor("swarmpal_fac")
class PalFacDataTreeAccessor:
def __init__(self, datatree) -> None:
self._datatree = datatree
[docs] def quicklook(self, active_tree="."):
fig, axes = plt.subplots(nrows=2, sharex=True)
# TODO: refactor to be able to identify active tree
process_config = self._datatree.swarmpal.pal_meta[active_tree]["FAC_singlesat"]
measurement_varname = process_config.get("measurement_varname", "B_NEC")
model_varname = process_config.get("model_varname", "B_NEC_Model")
dataset = process_config.get("dataset")
residual = (
self._datatree[f"{active_tree}/{dataset}"][measurement_varname]
- self._datatree[f"{active_tree}/{dataset}"][model_varname]
)
residual.diff(dim="Timestamp").plot.line(ax=axes[0], hue="NEC")
self._datatree[f"{active_tree}/{dataset}/output"]["FAC"].plot.line(ax=axes[1])
axes[0].set_xlabel("")
axes[0].set_ylabel(r"d/dt($\Delta$B)")
axes[0].grid()
axes[1].grid()
return fig, axes