Repository: RobertJaro/NF2 (language: python)

Issue:
increase default epochs



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index c8f15ae..fb66532 100644
--- a/README.md
+++ b/README.md
@@ -18,124 +18,213 @@ A systematic comparison of the time evolution of free magnetic energy and magnet
 
 NF2 can be used as framework to download SHARP data, perform extrapolations and to verify the results. 
 The colab notebook provides a basic example that can be adjusted for arbitrary active regions.
+For local usage configuration files can be used to perform extrapolations of single active regions or series.
 
-### Scripts (command line)
-
-Before you start download and install the requirements.txt in your Python3 environment.
-Make sure that PyTorch is properly installed if you use a GPU environment.
-
-```
-pip install git+https://github.com/RobertJaro/NF2.git
-```
-
-In the config file you have to specify the path to your data. You need the three vector components (Bp, Br, Bt) and their corresponding errors.
-The base path defines where the results and the trained model are stored. The bin defines spatial reduction of the input data (bin 2 is recommended for SHARP data).
-Other parameters are used for model training and can be set to their default value.
-
-Example for SHARP 377 (hmi_7115.json):
-```json
-{
-  "base_path": "/<<your path>>/hmi_377",
-  "logging": {
-    "wandb_entity": "<<your wandb username>>",
-    "wandb_project": "<<your project name>>",
-    "wandb_name": "hmi_377",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "data_path": [
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp.fits",
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp_err.fits",
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt.fits",
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt_err.fits",
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br.fits",
-    "/<<your path>>/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br_err.fits"
-    ],
-    "bin": 2,
-    "height_mapping": {"z":  [0.000]},
-    "Mm_per_pixel": 0.72,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 160,
-    "b_norm": 2500,
-    "spatial_norm": 160,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": null,
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
+## Installation
+
+Use pip to install the NF2 package.
+```
+pip install nf2
 ```
 
-Training requires about 1 hour on a single V100 GPU.
+For the latest version use the following command.
+```
+pip install git+https://github.com/RobertJaro/NF2@v0.3.0
+```
 
+Make sure that PyTorch is installed if you use a GPU environment. Run in python console:
+```python
+import torch
+    
+print('PyTorch version:', torch.__version__)
+print('GPU available:', torch.cuda.is_available())
+print('Number of GPUs:', torch.cuda.device_count())
 ```
-python -m nf2.extrapolate --config <<your path>>/hmi_7115.json
+
+## Scripts (command line)
+
+NF2 runs can be configured through yaml files.
+The method can be used for single active regions or time series.
+
+### Data download
+
+NF2 provides direct data download of SHARP data from JSOC.
+
+Download of single SHARP region:
+```
+nf2-download --download_dir "<<PATH TO SAVE>>" --email <<YOUR JSOC REGISTERED EMAIL>> --noaa_num 11158 --t_start 2011-02-15T00:00:00
 ```
 
-For the extrapolations of time series we can use the model weights of the previous time step. A single time step can then be performed in a few minutes.
-For this we also create a configuration that specifies our data.
-
-Example for the time series of SHARP 377 (hmi_377_series.json):
-```json
-{
-  "base_path": "/<<your path>>/hmi_series_377",
-  "meta_path": "/<<your previous training path>>/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "<<your wandb username>>",
-    "wandb_project": "<<your project name>>",
-    "wandb_name": "hmi_series_377",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/<<your path>>/",
-    "bin": 2,
-    "height_mapping": {"z":  [0.000]},
-    "Mm_per_pixel": 0.72,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 160,
-    "b_norm": 2500,
-    "spatial_norm": 160,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 2e3,
-    "work_directory": null,
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
+Download of SHARP series:
+```
+nf2-download --download_dir "<<PATH TO SAVE>>" --email <<YOUR JSOC REGISTERED EMAIL>> --noaa_num 11158 --t_start 2011-02-15T00:00:00 --t_end 2011-02-16T00:00:00
+```
+
+To use near real-time data add `--series sharp_cea_720s_nrt`
+
+### AR - Extrapolation
+The code provides NLFF extrapolations where a trade-off between the observation and the force-free magnetic field assumption is found. The weighting is controlled through the lambda parameters in the configuration file. Increase the force-free lambda to enforce the force-free magnetic field assumption. If solutions closer to the observation are desired, decrease the force-free lambda.
+
+Basic example for NOAA AR 11158 (377.yaml):
+```yaml
+---
+base_path: "<<PATH TO SAVE THE RESULTS>>"
+work_directory: "<<OPTIONAL - SCRATCH DIRECTORY>>"
+logging:
+  project: "<<YOUR WANDB PROJECT>>"
+  name: "SHARP 377"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br.fits"
+        Bt: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt.fits"
+        Bp: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp.fits"
+      error_path: # optional
+        Br_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: b
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
+```
+The training can be started with the following command and requires about 1 hour on a single V100 GPU.
+```
+nf2-extrapolate --config <<PATH TO CONFIG FILE>>/377.yaml
+```
+
+### AR - Divergence free
+
+To enforce a divergence free magnetic field, the magnetic field can be computed through a vector potential.
+For this we specify the model type as vector_potential and remove the divergence loss term. The divergence free condition can be kept for monitoring but should be close to zero. 
+Extrapolations that use the vector potential will require more computational resources and training time since the optimization is performed through second order derivatives.
+
+Example for a divergence-free extrapolation of NOAA AR 11158 (377_vp.yaml):
+```yaml
+---
+base_path: "<<PATH TO SAVE THE RESULTS>>"
+work_directory: "<<OPTIONAL - SCRATCH DIRECTORY>>"
+logging:
+  project: "<<YOUR WANDB PROJECT>>"
+  name: "SHARP 377 - VP"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br.fits"
+        Bt: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt.fits"
+        Bp: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp.fits"
+      error_path: # optional
+        Br_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "<<ABSOLUTE PATH TO DATA>>/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
 ```
 
+```
+nf2-extrapolate --config <<PATH TO CONFIG FILE>>/377_vp.yaml
+```
 
+### Time series
+
+For the extrapolations of time series we can use the model weights of the previous time step. A single time step can then be performed in a few minutes.
+For this we also create a configuration that specifies our data set.
+
+Example for the time series of NOAA 13664 (related to the May 2024 geomagnetic storm; SHARP 11149).
+
+(initial frame - 13664.yaml)
+```yaml
+---
+base_path: "<<PATH TO SAVE THE RESULTS>>/init"
+work_directory: "<<OPTIONAL - SCRATCH DIRECTORY>>"
+logging:
+  project: "<<YOUR WANDB PROJECT>>"
+  name: "NOAA 13664 - init"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Br.fits"
+        Bt: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bt.fits"
+        Bp: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Br_err.fits"
+        Bt_err: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bt_err.fits"
+        Bp_err: "<<ABSOLUTE PATH TO DATA>>//hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
+```
+(series - 13664_series.yaml)
+
+```yaml
+---
+base_path: "<<PATH TO SAVE THE RESULTS>>/series"
+work_directory: "<<OPTIONAL - SCRATCH DIRECTORY>>"
+meta_path: "<<PATH TO SAVE THE RESULTS>>/init/last.ckpt"
+logging:
+  project: "<<YOUR WANDB PROJECT>>"
+  name: "NOAA 13664 - series"
+data:
+  type: sharp
+  data_path: "<<ABSOLUTE PATH TO DATA>>"
+  num_workers: 8
+  iterations: 2.0e+3
+model:
+  type: vector_potential
+  dim: 256
+training:
+  check_val_every_n_epoch: 5 # validation plots in 1h steps
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
+```
 
+Run initial extrapolation and series extrapolation:
 ``` 
-python -m nf2.extrapolate_series --config <<your path>>/hmi_377_series.json
+nf2-extrapolate --config <<your path>>/13664.yaml
+nf2-extrapolate-series --config <<your path>>/13664_series.yaml
 ```
 
 ### Visualization
@@ -145,17 +234,66 @@ https://www.paraview.org/download/
 
 ![](images/paraview.jpeg)
 
-NF2 models can be converted to VTK files that can be used by Paraview.
+NF2 models can be converted to VTK files that can be used by Paraview. 
+Full resolution extrapolations typically exceed the memory capacity. 
+`--Mm_per_pixel 0.72` can be used to reduce the resolution to bin2 SHARP data (the default SHARP resolution is 0.36 Mm per pixel).
 
 ``` 
-python -m nf2.evaluation.convert_vtk_file <<path to your nf2 file>> <<path to the output vtk file>>
+nf2-to-vtk --nf2_path <<path to your nf2 file>> --out_path <<path to the output vtk file>> --Mm_per_pixel 0.72
 ```
 
-## Paper
+## Publications
+
+### Original method paper
+[Jarolim et al. 2023, Nature Astronomy](https://doi.org/10.1038/s41550-023-02030-9) \
+Jarolim, R., Thalmann, J.K., Veronig, A.M. et al. **Probing the solar coronal magnetic field with physics-informed neural networks**. Nat Astron 7, 1171–1179 (2023). https://doi-org.cuucar.idm.oclc.org/10.1038/s41550-023-02030-9
+
+BibTeX:
+```
+@ARTICLE{2023NatAs...7.1171J,
+       author = {{Jarolim}, R. and {Thalmann}, J.~K. and {Veronig}, A.~M. and {Podladchikova}, T.},
+        title = "{Probing the solar coronal magnetic field with physics-informed neural networks.}",
+      journal = {Nature Astronomy},
+         year = 2023,
+        month = oct,
+       volume = {7},
+        pages = {1171-1179},
+          doi = {10.1038/s41550-023-02030-9},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2023NatAs...7.1171J},
+      adsnote = {Provided by the SAO/NASA Astrophysics Data System}
+}
+```
+
+### Multi-height extrapolations
+[Jarolim et al. 2024, ApJL](https://doi.org/10.3847/2041-8213/ad2450) \
+Jarolim, R., Tremblay, B., Rempel, M., Molnar, M., Veronig, A. M., Thalmann, J. K., & Podladchikova, T. (2024). **Advancing solar magnetic field extrapolations through multiheight magnetic field measurements**. The Astrophysical Journal Letters, 963(1), L21.
+
+BibTeX:
+```
+@ARTICLE{2024ApJ...963L..21J,
+       author = {{Jarolim}, Robert and {Tremblay}, Benoit and {Rempel}, Matthias and {Molnar}, Momchil and {Veronig}, Astrid M. and {Thalmann}, Julia K. and {Podladchikova}, Tatiana},
+        title = "{Advancing Solar Magnetic Field Extrapolations through Multiheight Magnetic Field Measurements}",
+      journal = {\apjl},
+     keywords = {Solar magnetic fields, Neural networks, Solar corona, 1503, 1933, 1483},
+         year = 2024,
+        month = mar,
+       volume = {963},
+       number = {1},
+          eid = {L21},
+        pages = {L21},
+          doi = {10.3847/2041-8213/ad2450},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2024ApJ...963L..21J},
+      adsnote = {Provided by the SAO/NASA Astrophysics Data System}
+}
+```
+
+### Applications
+
+Purkhart, S., Veronig, A.M., Dickson, E.C., Battaglia, A.F., Krucker, S., Jarolim, R., Kliem, B., Dissauer, K. and Podladchikova, T., 2023. **Multipoint study of the energy release and transport in the 28 March 2022, M4 flare using STIX, EUI, and AIA during the first Solar Orbiter nominal mission perihelion**. Astronomy & Astrophysics, 679, p.A99.
 
-accepted; Nature Astronomy (Jarolim et al. 2023)
+McKevitt, J., Jarolim, R., Matthews, S., Baker, D., Temmer, M., Veronig, A., Reid, H. and Green, L., 2024. **The Link between Nonthermal Velocity and Free Magnetic Energy in Solar Flares**. The Astrophysical Journal Letters, 961(2), p.L29.
 
-Preprint available: https://doi.org/10.21203/rs.3.rs-1415262/v1
+Korsós, M.B., Jarolim, R., Erdélyi, R., Veronig, A.M., Morgan, H. and Zuccarello, F., 2024. **First Insights into the Applicability and Importance of Different 3D Magnetic Field Extrapolation Approaches for Studying the Preeruptive Conditions of Solar Active Regions**. The Astrophysical Journal, 962(2), p.171.
 
 ## Data
 All our simulation results are publicly available (parameter variation, time series, 66 individual active regions).
diff --git a/applications/AR13664.md b/applications/AR13664.md
new file mode 100644
index 0000000..c2241a2
--- /dev/null
+++ b/applications/AR13664.md
@@ -0,0 +1,50 @@
+# AR 13664
+
+On 2024 May 10/11, the strongest geomagnetic storm since November 2003 has occurred,
+with a peak Dst index of $-412$~nT. The storm was caused by NOAA Active Region (AR) 13664,
+which was the source of a large number of fast coronal mass ejections (CMEs) and big flares,
+including 12 X-class flares.
+
+We perform 3D magnetic field extrapolations with the **NF2** nonlinear-force free code based on physics
+informed neural networks.
+This tutorial provides an overview on how to access the derived data products and how to load snapshots for further
+analysis.
+
+**This page is currently under construction**
+**The publication is currently under review**
+
+## Overview
+
+<img src="images/AR13664_overview.jpg">
+
+Overview of the evolution of NOAA AR 13664 from 2024 May 5-00:00 to 11-04:36 UT. a) Evolution of the vertical
+distribution of free magnetic energy. b)
+GOES 1-8~{\AA} soft X-ray flux. c) Observed SDO/HMI radial magnetic field maps, SDO/AIA~131~{\AA} EUV maps, and the
+modeled current density integrated along the vertical axis. We show the initial flux emergence (snapshot at 2024 May 7
+00:00~UT) and the major flare eruptions (indicated by dashed lines). The magenta line indicates the point where the
+right edge of the SHARP reaches 60$^\circ$ longitude, beyond which the reliability of the vector magnetograms and the
+modeled coronal field decreases. The circles in the integrated current density maps indicate points of flux domain
+interactions. An animation of the full time series is available in the online journal. Left top: B$_z$ of the modeled
+photospheric magnetic field. Right top: EUV observations from SDO/AIA~131~{\AA} EUV. Left bottom: maps of integrated
+current density. Right bottom: maps of integrated free magnetic energy. The sequence shows the re-configuration of
+magnetic domains (dark lines in current density maps) during strong solar eruptions.
+
+## Data archive
+
+All data products are publicly available through [Globus](https://app.globus.org/file-manager?origin_id=4263de78-cfdb-401e-a62b-dae3b935530a&origin_path=%2F) website.
+
+## Load hdf5 snapshot
+
+**in progress**
+
+## Load NF2 snapshot
+
+**in progress**
+
+## Load evaluation results
+
+**in progress**
+
+## Publication
+
+**in progress**
\ No newline at end of file
diff --git a/applications/images/AR13664_overview.jpg b/applications/images/AR13664_overview.jpg
new file mode 100644
index 0000000..66a94d6
Binary files /dev/null and b/applications/images/AR13664_overview.jpg differ
diff --git a/config/2106_AR.json b/config/2106_AR.json
new file mode 100644
index 0000000..7de5c8e
--- /dev/null
+++ b/config/2106_AR.json
@@ -0,0 +1,78 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2106_AR_v02",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2106_AR_v02",
+  "logging": {
+    "project": "full_disc_v2",
+    "name": "2106"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "AR", "batch_size": 24576,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {
+            "type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [-0.86, 1.65], "invert": true
+          }
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random",
+          "latitude_range": [0.47, 2.67], "longitude_range": [-0.86, 1.65]
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [-0.86, 1.65], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "n_slices": 10, "plot_currents": true, "latitude_range": [0.47, 2.67], "longitude_range": [-0.86, 1.65]
+        }
+      ],
+      "max_radius": 1.3,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 2e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  10,
+        "ds_id": "AR"
+      },
+      {"type": "force_free", "lambda": 1e-1},
+      {"type":  "divergence", "lambda":  1e-1},
+      {"type":  "energy_gradient", "lambda":  1e-1, "base_radius":  1.0}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/2173.json b/config/2173.json
new file mode 100644
index 0000000..631155f
--- /dev/null
+++ b/config/2173.json
@@ -0,0 +1,86 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2173_synoptic_v08",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_synoptic_v08",
+  "logging": {
+    "project": "subframe_v2",
+    "name": "vector potential"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+      {
+          "type": "map", "name": "synoptic", "batch_size": 16384,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          }
+        },
+        {
+          "type": "random_spherical", "batch_size": 32768, "name": "random"
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "synoptic_valid", "filter_nans": false, "shuffle": false, "strides": 4, "plot_overview": false,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          }
+        },
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "n_slices": 10, "plot_currents": true
+        },
+        {
+          "type": "spherical_slices", "name": "subframe_slices",
+          "n_slices": 10, "plot_currents": true,
+          "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559]
+        },
+        {
+          "type": "pfss_boundary", "name": "pfss_boundary_valid", "shuffle": false,  "plot_overview": false,
+          "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits"
+        }
+      ],
+      "max_radius": 1.3,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "vector_potential",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  10,
+        "ds_id":  "synoptic"
+      },
+      {"type": "force_free", "lambda": 1e-1},
+      {"type":  "energy_gradient", "lambda":  1e-2, "base_radius":  1.0}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/2173_full_disc.json b/config/2173_full_disc.json
new file mode 100644
index 0000000..4bc2878
--- /dev/null
+++ b/config/2173_full_disc.json
@@ -0,0 +1,81 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2173_fulldisc_v01",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_fulldisc_v01",
+  "logging": {
+    "project": "full_disc_v2",
+    "name": "2173"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "AR", "batch_size": 24576,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [2, 5], "invert": true}
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random",
+          "latitude_range": [0.47, 2.67], "longitude_range": [2, 5]
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [2, 5], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "n_slices": 10, "plot_currents": true, "latitude_range": [0.47, 2.67], "longitude_range": [2, 5]
+        },
+        {
+          "type": "spherical_slices", "name": "subframe_slices",
+          "n_slices": 10, "plot_currents": true,
+          "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559]
+        }
+      ],
+      "max_radius": 1.3,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  10,
+        "ds_id": "AR"
+      },
+      {"type": "force_free", "lambda": 1e-1},
+      {"type":  "divergence", "lambda":  1e-1},
+      {"type":  "energy_gradient", "lambda":  1e-1, "base_radius":  1.0}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/2173_submap.json b/config/2173_submap.json
new file mode 100644
index 0000000..9a8ab56
--- /dev/null
+++ b/config/2173_submap.json
@@ -0,0 +1,101 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2173_v04",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_v04",
+  "logging": {
+    "project": "subframe_v2",
+    "name": "2173"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "AR", "batch_size": 24576,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+          }
+        },
+        {
+          "type": "map", "name": "synoptic", "batch_size": 8192,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          },
+          "mask_configs" : {
+            "type":  "reference",
+            "file": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+          }
+        },
+        {
+          "type": "pfss_boundary", "name": "pfss_boundary", "batch_size": 8192,
+          "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits"
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random"
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "synoptic_valid", "filter_nans": false, "shuffle": false, "strides": 4, "plot_overview": false,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          }
+        },
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "n_slices": 10, "plot_currents": true
+        },
+        {
+          "type": "spherical_slices", "name": "subframe_slices",
+          "n_slices": 10, "plot_currents": true,
+          "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559]
+        }
+      ],
+      "max_radius": 1.5,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  {"start":  1e3, "end":  1e1, "iterations": 1e5},
+        "ds_id": ["synoptic", "AR", "pfss_boundary"]
+      },
+      {"type": "force_free", "lambda": 1e-1},
+      {"type":  "divergence", "lambda":  1e-1}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/2173_synoptic_ref.json b/config/2173_synoptic_ref.json
new file mode 100644
index 0000000..bceefb6
--- /dev/null
+++ b/config/2173_synoptic_ref.json
@@ -0,0 +1,75 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2173_fulldisc_v02",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_fulldisc_v02",
+  "logging": {
+    "project": "full_disc_v2",
+    "name": "2173"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "AR", "batch_size": 24576,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [2, 5], "invert": true}
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random",
+          "latitude_range": [0.47, 2.67], "longitude_range": [2, 5]
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "latitude_range": [0.47, 2.67], "longitude_range": [2, 5], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "n_slices": 10, "plot_currents": true, "latitude_range": [0.47, 2.67], "longitude_range": [2, 5]
+        },
+        {
+          "type": "spherical_slices", "name": "subframe_slices",
+          "n_slices": 10, "plot_currents": true,
+          "longitude_range": [2.35619449, 3.14159265], "latitude_range": [1.25663706, 1.88495559]
+        }
+      ],
+      "max_radius": 1.5,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  10,
+        "ds_id": "AR"
+      },
+      {"type": "force_free", "lambda": 1e-1},
+      {"type":  "divergence", "lambda":  1e-1},
+      {"type":  "energy_gradient", "lambda":  1e-1, "base_radius":  1.0}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/377_az.json b/config/377_az.json
new file mode 100644
index 0000000..a860981
--- /dev/null
+++ b/config/377_az.json
@@ -0,0 +1,58 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/disambig/377_az_v4",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "disambiguation",
+    "wandb_name": "377_az",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "azimuth",
+    "B_data": {
+        "B_field": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.field.fits",
+        "B_inclination": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.inclination.fits",
+        "B_azimuth": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.azimuth.fits",
+        "B_disambig": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.disambig.fits"
+      },
+    "slice": {
+        "longitude": [0.3, 0.9],
+        "latitude": [1.0, 1.4],
+        "frame": "heliographic_carrington"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.2,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 8192,
+      "random": 16384
+    },
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/377_az_v4",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
+    "lambda_force_free": 1e-2,
+    "lambda_energy_gradient": 1e-2,
+    "lambda_potential": 1e-3,
+    "check_val_every_n_epoch": 10,
+    "boundary_loss": "azimuth"
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "frame",
+      "longitude_range": [0.39223768, 0.808977],
+      "latitude_range": [1.11188089, 1.30852593],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/analytical_0_8_16_32.json b/config/analytical_0_8_16_32.json
index ab913be..b5ddb90 100644
--- a/config/analytical_0_8_16_32.json
+++ b/config/analytical_0_8_16_32.json
@@ -1,5 +1,5 @@
 {
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_0_8_16_32",
+  "base_path": "/glade/work/rjarolim/nf2/analytical/analytical_0_8_16_32",
   "logging": {
     "wandb_entity": "robert_jarolim",
     "wandb_project": "multi-height",
@@ -16,7 +16,7 @@
     "spatial_norm": 64,
     "batch_size": {"boundary":  1e4, "random":  2e4},
     "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_0_8_16_32",
+    "work_directory": "/glade/work/rjarolim/nf2/analytical/analytical_0_8_16_32",
     "num_workers": 8,
     "Mm_per_pixel": 1
   },
diff --git a/config/demo_submap.json b/config/demo_submap.json
new file mode 100644
index 0000000..6274b44
--- /dev/null
+++ b/config/demo_submap.json
@@ -0,0 +1,103 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2106_v03",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2106_v03",
+  "logging": {
+    "project": "subframe_v2",
+    "name": "2106"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "AR", "batch_size": 32768,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "longitude_range": [0.3, 0.9], "latitude_range": [1.0, 1.4], "invert": true
+          }
+        },
+        {
+          "type": "random_spherical", "batch_size": 32768, "name": "random_subframe",
+          "longitude_range": [0.3, 0.9], "latitude_range": [1.0, 1.4], "radius_range": [1, 1.2]
+        },
+        {
+          "type": "map", "name": "synoptic", "batch_size": 16384,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "longitude_range": [0.3, 0.9], "latitude_range": [1.0, 1.4], "invert": false}
+        },
+        {
+          "type": "random_spherical", "batch_size": 16384, "name": "random"
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "synoptic_valid", "filter_nans": false, "shuffle": false, "strides": 4, "plot_overview": false,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits"
+          }
+        },
+        {
+          "type": "map", "name": "AR_valid", "filter_nans": false, "shuffle": false, "plot_overview": false,
+          "files" : {
+            "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+            "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+            "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+            "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+            "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+            "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+          },
+          "mask_configs" : {"type":  "heliographic_carrington",
+            "longitude_range": [0.3, 0.9], "latitude_range": [1.0, 1.4], "invert": true}
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "radius_range": [1.0, 1.3], "n_slices": 5, "plot_currents": true
+        },
+        {
+          "type": "spherical_slices", "name": "subframe_slices",
+          "radius_range": [1.0, 1.3], "n_slices": 10, "plot_currents": true,
+          "longitude_range": [0.3, 0.9], "latitude_range": [1.0, 1.4]
+        }
+      ],
+      "max_radius": 2.5,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  {"start":  1e3, "end":  1e1, "iterations": 1e5},
+        "ds_id": ["synoptic", "AR"]
+      },
+      {"type": "force_free", "lambda": 1},
+      {"type":  "divergence", "lambda": 1},
+      {"type":  "radial", "lambda": 1e-2}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/demo_synoptic.json b/config/demo_synoptic.json
new file mode 100644
index 0000000..3c72ec4
--- /dev/null
+++ b/config/demo_synoptic.json
@@ -0,0 +1,66 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/2106_synoptic_v07",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2106_synoptic_v07",
+  "logging": {
+    "project": "synoptic_v2",
+    "name": "2106"
+},
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "synoptic", "batch_size": 65536,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits"
+          }
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random",
+          "radial_weighted_sampling": true, "latitude_weighted_sampling": true
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "synoptic_valid", "filter_nans": false, "shuffle": false, "strides": 4, "plot_overview": false,
+          "files" : {
+            "Br": "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits",
+            "Bp": "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits"
+          }
+        },
+        {
+          "type": "sphere", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk",
+          "radius_range": [1.0, 2.5], "n_slices": 10, "plot_currents": true
+        }
+      ],
+      "max_radius": 2.5,
+      "G_per_dB": 2000,
+      "num_workers": 8,
+      "batch_size": 65536
+  },
+  "model": {
+    "type": "b",
+    "dim": 512,
+    "activation": "swish"
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {
+        "type": "boundary", "name":  "boundary",
+        "lambda":  {"start":  1e3, "end":  1e1, "iterations": 1e5},
+        "ds_id": "synoptic"
+      },
+      {"type": "force_free", "lambda": 1},
+      {"type":  "divergence", "lambda": 1},
+      {"type":  "radial", "lambda": 1e-2}
+    ],
+    "check_val_every_n_epoch": 25
+  }
+}
\ No newline at end of file
diff --git a/config/fd_2106.json b/config/fd_2106.json
new file mode 100644
index 0000000..78f1794
--- /dev/null
+++ b/config/fd_2106.json
@@ -0,0 +1,69 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/full_disk/2106_v1",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "full_disk",
+    "wandb_name": "2106",
+    "wandb_id": null
+  },
+  
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.b_synoptic.2106.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+      "Bt":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.b_synoptic.2106.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.b_720s.20110215_000000_TAI.Br.fits",
+      "Bt":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2106/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+    },
+    "boundary": {
+      "type":  "open"
+    },
+    "height": 1.2,
+    "b_norm": 2000,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/tmp/nf2/2106_v1",
+    "num_workers": 8
+  },
+
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": false,
+    "use_positional_encoding": false
+  },
+
+  "training": {
+    "epochs": 1000,
+    "lambda_b": {"start": 1e3, "end": 10, "iterations" : 2e5},
+    "lambda_div": 1e-2,
+    "lambda_ff": 1e-2,
+    "lambda_min_energy": 0,
+    "lambda_radial_reg": 0,
+    "check_val_every_n_epoch": 5,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 2e5}
+  },
+
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [0.39223768, 0.808977],
+      "latitude_range": [1.11188089, 1.30852593],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2106_subframe_vp.json b/config/fd_2106_subframe_vp.json
new file mode 100644
index 0000000..98e60c2
--- /dev/null
+++ b/config/fd_2106_subframe_vp.json
@@ -0,0 +1,71 @@
+{
+  "base_path": "/glade/work/rjarolim/global/subframe/2106_subframe_vp_v1",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_subframe",
+    "wandb_name": "2106_subframe_vp",
+    "wandb_id": null
+  },
+
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp":  "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits",
+      "Br":  "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+      "Bt":  "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+      "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+      "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+      "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+      "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+      "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.2,
+    "slice": {
+      "longitude": [0.3, 0.9],
+      "latitude": [1.0, 1.4],
+      "frame":  "heliographic_carrington"
+    },
+    "b_norm": 2000,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2106_subframe_vp_v1",
+    "num_workers": 8
+  },
+
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": true,
+    "use_positional_encoding": false
+  },
+
+  "training": {
+    "epochs": 1e6,
+    "lambda_b": {"start": 1e2, "end": 10, "iterations" : 1e5},
+    "lambda_divergence": 0,
+    "lambda_force_free": 1e-1,
+    "lambda_min_energy": 0,
+    "lambda_radial_reg": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 20,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "iterations": 1e5}
+  },
+
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [0.39223768, 0.808977],
+      "latitude_range": [1.11188089, 1.30852593],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2106_vp.json b/config/fd_2106_vp.json
new file mode 100644
index 0000000..06468fc
--- /dev/null
+++ b/config/fd_2106_vp.json
@@ -0,0 +1,73 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/global/2106_vp_v2",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2106_vp",
+    "wandb_id": null
+  },
+  
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp":  "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bp.fits",
+      "Br":  "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits",
+      "Bt":  "/glade/work/rjarolim/data/global/fd_2106/hmi.b_synoptic.2106.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits",
+      "Br":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits",
+      "Bt":  "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits",
+      "Bp_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits",
+      "Br_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits",
+      "Bt_err": "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type":  "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2106_vp_v2",
+    "num_workers": 8
+  },
+
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": true,
+    "use_positional_encoding": false
+  },
+
+  "training": {
+    "epochs": 1e6,
+    "lambda_b": {"start": 1e2, "end": 10, "iterations" : 1e5},
+    "lambda_divergence": 0,
+    "lambda_force_free": 1e-1,
+    "lambda_min_energy": 0,
+    "lambda_radial_reg": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 5,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "iterations": 1e5}
+  },
+
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [0.39223768, 0.808977],
+      "latitude_range": [1.11188089, 1.30852593],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.2],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2154.json b/config/fd_2154.json
new file mode 100644
index 0000000..9b2b8c6
--- /dev/null
+++ b/config/fd_2154.json
@@ -0,0 +1,63 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/full_disk/2154_v4",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global",
+    "wandb_name": "2154",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_synoptic.2154.Bp.fits",
+      "Br": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.synoptic_mr_polfil_720s.2154.Mr_polfil.fits",
+      "Bt": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_synoptic.2154.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bp.fits",
+      "Br": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Br.fits",
+      "Bt": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bt.fits",
+      "Bp_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bp_err.fits",
+      "Br_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Br_err.fits",
+      "Bt_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 2.5,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 8192,
+      "random": 16384
+    },
+    "work_directory": "/tmp/nf2/2154_v4",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": false
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-2,
+    "lambda_divergence": 1e-2,
+    "lambda_radial": 0,
+    "lambda_potential": 0,
+    "lambda_energy_gradient": 0,
+    "check_val_every_n_epoch": 1
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 2.0],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2154_az.json b/config/fd_2154_az.json
new file mode 100644
index 0000000..9664e86
--- /dev/null
+++ b/config/fd_2154_az.json
@@ -0,0 +1,62 @@
+{
+  "base_path": "/glade/work/rjarolim/global/2154_az_v1",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2154_az",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "azimuth",
+    "files": {
+      "B_field": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.field.fits",
+      "B_inclination": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.inclination.fits",
+      "B_azimuth": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.azimuth.fits",
+      "B_disambig": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.B_720s.20140902_060000_TAI.disambig.fits"
+    },
+    "auto_disambig": true,
+    "boundary": {
+      "type": "potential",
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits"
+    },
+    "height": 1.5,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 8192,
+      "random": 16384
+    },
+    "work_directory": "/tmp/nf2/2154_az_v1",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": false,
+    "use_positional_encoding": false
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_ff": 1e-2,
+    "check_val_every_n_epoch": 100,
+    "lr_params": {
+      "start": 5e-4,
+      "end": 5e-5,
+      "decay_iterations": 2e5
+    }
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.5],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2154_vp.json b/config/fd_2154_vp.json
new file mode 100644
index 0000000..c6e044a
--- /dev/null
+++ b/config/fd_2154_vp.json
@@ -0,0 +1,61 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/global/2154_vp_v11",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2154_vp",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2154/hmi.b_synoptic.2154.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2154/hmi.b_synoptic.2154.Br.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2154/hmi.b_synoptic.2154.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Br.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bt.fits",
+      "Bp_err": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bp_err.fits",
+      "Br_err": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Br_err.fits",
+      "Bt_err": "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 2.0,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 32768,
+      "random": 65536
+    },
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2154_vp_v11",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-1,
+    "lambda_radial": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 5
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 2.0],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2154_vp_series.json b/config/fd_2154_vp_series.json
new file mode 100644
index 0000000..748da21
--- /dev/null
+++ b/config/fd_2154_vp_series.json
@@ -0,0 +1,52 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/full_disk/2154_vp_series_v1",
+  "meta_path": "/gpfs/gpfs0/robert.jarolim/full_disk/2154_vp_v6/extrapolation_result.nf2",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global",
+    "wandb_name": "2154_vp_series",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_synoptic.2154.Bp.fits",
+      "Br": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.synoptic_mr_polfil_720s.2154.Mr_polfil.fits",
+      "Bt": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_synoptic.2154.Bt.fits"
+    },
+    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk",
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 8192,
+      "random": 16384
+    },
+    "work_directory": "/tmp/nf2/2154_vp_series_v1",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "lambda_b": 10,
+    "lambda_ff": 1e-2,
+    "lambda_radial": 0,
+    "lambda_energy_gradient": 1e-3,
+    "check_val_every_n_epoch": 5,
+    "lr_params": 5e-4,
+    "reload_dataloaders_every_n_epochs": 5
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2173.json b/config/fd_2173.json
new file mode 100644
index 0000000..88fd765
--- /dev/null
+++ b/config/fd_2173.json
@@ -0,0 +1,73 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/full_disk/2173_v3",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "full_disk",
+    "wandb_name": "2173",
+    "wandb_id": null
+  },
+  
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_synoptic.2173.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+      "Bt":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_synoptic.2173.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Br.fits",
+      "Bt": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Bt.fits",
+      "Bp_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Bp_err.fits",
+      "Br_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Br_err.fits",
+      "Bt_err": "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2173/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.5,
+    "b_norm": 2000,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/tmp/nf2/2173_v3",
+    "num_workers": 8
+  },
+
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": false,
+    "use_positional_encoding": false
+  },
+
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {"start": 1e2, "end": 10, "iterations" : 2e5},
+    "lambda_div": 1e-2,
+    "lambda_ff": 1e-2,
+    "lambda_min_energy": 0,
+    "lambda_radial_reg": 0,
+    "lambda_energy_gradient_reg": 1e-2,
+    "check_val_every_n_epoch": 5,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 2e5}
+  },
+
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [2.356194490192345, 3.141592653589793],
+      "latitude_range": [1.2566370614359172, 1.8849555921538759],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 10,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.5],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2173_synoptic.json b/config/fd_2173_synoptic.json
new file mode 100644
index 0000000..531e698
--- /dev/null
+++ b/config/fd_2173_synoptic.json
@@ -0,0 +1,62 @@
+{
+  "base_path": "/glade/work/rjarolim/global/2173_synoptic_v1",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2173_synoptic",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "synoptic",
+    "synoptic_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 32768,
+      "random": 65536
+    },
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_synoptic_v1",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-2,
+    "lambda_radial": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 5
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [2.356194490192345, 3.141592653589793],
+      "latitude_range": [1.2566370614359172, 1.8849555921538759],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 10,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2173_vp.json b/config/fd_2173_vp.json
new file mode 100644
index 0000000..0243ab3
--- /dev/null
+++ b/config/fd_2173_vp.json
@@ -0,0 +1,67 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/global/2173_vp_v3",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2173_vp",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt.fits",
+      "Bp_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp_err.fits",
+      "Br_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br_err.fits",
+      "Bt_err": "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt_err.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_vp_v3",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e6,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-1,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 5,
+    "lr_params": 5e-4
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [2.356194490192345, 3.141592653589793],
+      "latitude_range": [1.2566370614359172, 1.8849555921538759],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 10,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2173_vp_series.json b/config/fd_2173_vp_series.json
new file mode 100644
index 0000000..660d310
--- /dev/null
+++ b/config/fd_2173_vp_series.json
@@ -0,0 +1,57 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/global/2173_vp_series_v2",
+  "meta_path": "/glade/work/rjarolim/nf2/global/2173_vp_v1/last.ckpt",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_series",
+    "wandb_name": "2173_vp_series",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "spherical",
+    "synoptic_files": {
+      "Bp": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits",
+      "Bt": "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits"
+    },
+    "paths": "/glade/work/rjarolim/data/global/fd_2173/full_disk",
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {"boundary": 8192, "random": 16384},
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_vp_series_v2",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "lambda_b": 10,
+    "lambda_force_free": 1e-1,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 5,
+    "lr_params":  5e-4,
+    "reload_dataloaders_every_n_epochs": 5
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [2.356194490192345, 3.141592653589793],
+      "latitude_range": [1.2566370614359172, 1.8849555921538759],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 10,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/fd_2213_synoptic.json b/config/fd_2213_synoptic.json
new file mode 100644
index 0000000..fcecefb
--- /dev/null
+++ b/config/fd_2213_synoptic.json
@@ -0,0 +1,60 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/synoptic/2213_v2",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_synoptic",
+    "wandb_name": "2213_synoptic",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "synoptic",
+    "synoptic_files": {
+      "Bp": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.2213.Bp.fits",
+      "Br": "/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.2213.Mr_polfil.fits",
+      "Bt": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.2213.Bt.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 2.5,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 32768,
+      "random": 65536
+    },
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2213_synoptic_v2",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-1,
+    "lambda_radial": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 10
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "outer-corona",
+      "radius_range": [1.0, 2.0],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/hmi_7115.json b/config/hmi_7115.json
index ab35fde..91004ee 100644
--- a/config/hmi_7115.json
+++ b/config/hmi_7115.json
@@ -1,5 +1,5 @@
 {
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115",
+  "base_path": "/Users/robert/PycharmProjects/NF2/test",
   "logging": {
     "wandb_entity": "robert_jarolim",
     "wandb_project": "multi-height",
@@ -10,15 +10,14 @@
   "data": {
     "type": "sharp",
     "data_path": [
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Bp.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Bp_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Bt.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Bt_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Br.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/7115/hmi.sharp_cea_720s.7115.20170906_000000_TAI.Br_err.fits"
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Bp.fits",
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Bp_err.fits",
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Bt.fits",
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Bt_err.fits",
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Br.fits",
+    "/Users/robert/PycharmProjects/NF2/test/hmi.sharp_cea_720s.7115.20170902_000000_TAI.Br_err.fits"
     ],
     "bin": 2,
-    "height_mapping": {"z":  [0.000]},
     "Mm_per_pixel": 0.72,
     "boundary": {"type":  "potential", "strides":  4},
     "height": 160,
@@ -31,8 +30,7 @@
   },
 
   "model": {
-    "dim": 256,
-    "use_height_mapping": false
+    "dim": 256
   },
 
   "training": {
diff --git a/config/hmi_7115_series.json b/config/hmi_7115_series.json
index 61f3b37..baf911b 100644
--- a/config/hmi_7115_series.json
+++ b/config/hmi_7115_series.json
@@ -1,18 +1,17 @@
 {
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_7115",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115/extrapolation_result.nf2",
+  "base_path": "<<YOUR PATH>>",
+  "meta_path": "<<Initial RUN>>/extrapolation_result.nf2",
   "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
+    "wandb_entity": "<<YOUR Entity>>",
+    "wandb_project": "nf2",
     "wandb_name": "hmi_series_7115",
     "wandb_id": null
   },
   
   "data": {
     "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/7115",
+    "paths": "<<DIR>>",
     "bin": 2,
-    "height_mapping": {"z":  [0.000]},
     "Mm_per_pixel": 0.72,
     "boundary": {"type":  "potential", "strides":  4},
     "height": 160,
@@ -25,8 +24,7 @@
   },
 
   "model": {
-    "dim": 256,
-    "use_height_mapping": false
+    "dim": 256
   },
 
   "training": {
diff --git a/config/spherical.json b/config/spherical.json
new file mode 100644
index 0000000..3f60470
--- /dev/null
+++ b/config/spherical.json
@@ -0,0 +1,47 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/spherical_512_vp",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "spherical",
+    "wandb_name": "2119",
+    "wandb_id": null
+  },
+  
+  "data": {
+    "type": "spherical",
+    "synchronic_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.b_synoptic.2119.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.b_synoptic.2119.Br.fits",
+      "Bt":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.b_synoptic.2119.Bt.fits"
+    },
+    "full_disk_files": {
+      "Bp":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.B_720s.20120123_033600_TAI.Bp.fits",
+      "Br":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.B_720s.20120123_033600_TAI.Br.fits",
+      "Bt":  "/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2012_01_23/hmi.B_720s.20120123_033600_TAI.Bt.fits"
+    },
+    "boundary": {"type":  "pfss"},
+    "height": 2.5,
+    "b_norm": 500,
+    "batch_size": {"boundary":  8192, "random":  16384},
+    "work_directory": "/tmp/nf2/spherical_vp",
+    "num_workers": 8
+  },
+
+  "model": {
+    "dim": 512,
+    "use_height_mapping": false,
+    "spherical": true,
+    "use_vector_potential": true
+  },
+
+  "training": {
+    "epochs": 1000,
+    "lambda_b": {"start": 1e2, "end": 1, "iterations" : 5e4},
+    "lambda_div": 0,
+    "lambda_ff": 1e-2,
+    "lambda_min_energy": 0,
+    "lambda_radial_reg": 0,
+    "check_val_every_n_epoch": 5,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
+  }
+}
\ No newline at end of file
diff --git a/config/sst.json b/config/sst.json
index d346600..7cd4336 100644
--- a/config/sst.json
+++ b/config/sst.json
@@ -1,51 +1,73 @@
 {
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/sst_v5",
+  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/sst_v1",
   "logging": {
     "wandb_entity": "robert_jarolim",
     "wandb_project": "multi-height",
     "wandb_name": "sst"
   },
-
   "data": {
     "type": "fits",
-    "data_paths": [
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_8542_2018-09-30_mosaic_Bvec_cea.fits"
-     ],
-    "mask_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea_mask.fits",
-    "height_mapping": {
-      "z":  [0.000, 10],
-      "z_min": [0.000, 0.000],
-      "z_max": [0.000, 200]
-    },
-    "flip_sign": 1,
-    "buffer": 0.5,
-    "Mm_per_pixel": 0.36,
+    "data": [
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Br_cea.fits"
+        },
+        "flip_sign": 1
+      },
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Br_cea.fits"
+        },
+        "flip_sign": 1
+      },
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Br_cea.fits"
+        },
+        "height_mapping": {
+          "z": 80.0,
+          "z_min": 0.0,
+          "z_max": 2500.0
+        },
+        "flip_sign": 1
+      }
+    ],
     "boundary": {
-      "type":  "open"
+      "type": "potential",
+      "index": 2,
+      "strides": 20
     },
     "b_norm": 2500,
-    "bin": 4,
-    "spatial_norm": 539.5,
-    "height": 539.5,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
+    "bin": 1,
+    "Mm_per_pixel": 0.036,
+    "spatial_norm": 3200,
+    "height": 3200,
+    "batch_size": {
+      "boundary": 1e4,
+      "random": 2e4
+    },
     "iterations": 1e5,
-    "work_directory": "/tmp/nf2/sst_v5",
+    "work_directory": "/tmp/nf2/sst_v1",
     "num_workers": 8,
-    "validation_strides": 2
+    "validation_strides": 10
   },
-
   "model": {
     "dim": 256,
     "use_height_mapping": true,
-    "use_vector_potential": true
+    "use_vector_potential": false
   },
-
   "training": {
+    "epochs": 100,
     "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_min_energy": 1e-3,
+    "lambda_div": 1e-3,
+    "lambda_ff": 1e-3,
+    "lambda_min_energy": 0,
     "lambda_height_reg": 1e-3,
     "validation_interval": 1e4,
     "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
diff --git a/config/sst_vp.json b/config/sst_vp.json
new file mode 100644
index 0000000..27fd1ed
--- /dev/null
+++ b/config/sst_vp.json
@@ -0,0 +1,75 @@
+{
+  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/sst_vp_v2",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "multi-height",
+    "wandb_name": "sst"
+  },
+  "data": {
+    "type": "fits",
+    "data": [
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_Br_cea.fits"
+        },
+        "flip_sign": 1
+      },
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_Br_cea.fits"
+        },
+        "flip_sign": 1
+      },
+      {
+        "B": {
+          "x": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Bp_cea.fits",
+          "y": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Bt_cea.fits",
+          "z": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_Br_cea.fits"
+        },
+        "height_mapping": {
+          "z": 80.0,
+          "z_min": 0.0,
+          "z_max": 2500.0
+        },
+        "flip_sign": 1
+      }
+    ],
+    "boundary": {
+      "type": "potential",
+      "index": 2,
+      "strides": 20
+    },
+    "b_norm": 2500,
+    "bin": 1,
+    "Mm_per_pixel": 0.036,
+    "spatial_norm": 3200,
+    "height": 3200,
+    "batch_size": {
+      "boundary": 1e4,
+      "random": 2e4
+    },
+    "iterations": 1e5,
+    "work_directory": "/tmp/nf2/sst_vp_v2",
+    "num_workers": 8,
+    "validation_strides": 10
+  },
+  "model": {
+    "dim": 256,
+    "use_height_mapping": true,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 10,
+    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
+    "lambda_div": 0,
+    "lambda_ff": 1e-3,
+    "lambda_min_energy": 0,
+    "lambda_height_reg": 0,
+    "validation_interval": 1e4,
+    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
+  }
+}
\ No newline at end of file
diff --git a/config/synoptic/synoptic_series.json b/config/synoptic/synoptic_series.json
new file mode 100644
index 0000000..1346944
--- /dev/null
+++ b/config/synoptic/synoptic_series.json
@@ -0,0 +1,47 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/synoptic/series_v3",
+  "meta_path": "/glade/work/rjarolim/nf2/synoptic/2213_v1/last.ckpt",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_synoptic",
+    "wandb_name": "series",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "synoptic",
+    "paths": {
+      "Bp": "/glade/work/rjarolim/data/global/synoptic/*.Bp.fits",
+      "Bt": "/glade/work/rjarolim/data/global/synoptic/*.Bt.fits",
+      "Br": "/glade/work/rjarolim/data/global/synoptic/*.Mr_polfil.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {"boundary": 8192, "random": 16384},
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/synoptic_series",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "lambda_b": 10,
+    "lambda_force_free": 1e-1,
+    "lambda_energy_gradient": 1e-2,
+    "lambda_radial": 0,
+    "check_val_every_n_epoch": 20,
+    "reload_dataloaders_every_n_epochs": 20
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file
diff --git a/config/synoptic/template.json b/config/synoptic/template.json
new file mode 100644
index 0000000..01fc600
--- /dev/null
+++ b/config/synoptic/template.json
@@ -0,0 +1,65 @@
+{
+  "base_path": "/glade/work/rjarolim/nf2/spherical/${carrington_rotation}_v01",
+  "work_directory": "/glade/derecho/scratch/rjarolim/nf2/${carrington_rotation}_synoptic_v01",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_synoptic",
+    "wandb_name": "${carrington_rotation}",
+    "wandb_id": null
+  },
+  "data":
+    {
+      "type": "spherical",
+      "train_configs": [
+        {
+          "type": "map", "name": "synoptic", "batch_size": 32768,
+          "files" : {
+            "Bp": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.${carrington_rotation}.Bp.fits",
+            "Br": "/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.${carrington_rotation}.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.${carrington_rotation}.Bt.fits"
+          }
+        },
+        {
+          "type": "random_spherical", "batch_size": 65536, "name": "random"
+        }
+      ],
+      "validation_configs": [
+        {
+          "type": "map", "name": "synoptic", "filter_nans": false,
+          "files" : {
+            "Bp": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.${carrington_rotation}.Bp.fits",
+            "Br": "/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.${carrington_rotation}.Mr_polfil.fits",
+            "Bt": "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.${carrington_rotation}.Bt.fits"
+          }
+        },
+        {
+          "type": "sphere", "set": "validation", "name": "sphere"
+        },
+        {
+          "type": "spherical_slices", "name": "full-disk", "set": "validation",
+          "radius_range": [1.0, 2.5], "n_slices": 5, "plot_currents": true
+        }
+      ],
+      "height": 1.3,
+      "b_norm": 2000,
+      "num_workers": 8
+  },
+  "model": {
+    "dim": 64,
+    "model": "flux",
+    "use_positional_encoding": true
+  },
+  "training": {
+    "epochs": 1e3,
+    "loss_config": [
+      {"type": "boundary", "lambda":  10},
+      {"type": "force_free", "lambda": 0},
+      {"type": "radial", "lambda": 0, "base_radius": 1.5},
+      {"type": "energy_gradient", "lambda": 0, "base_radius": 1.3},
+      {"type": "divergence", "lambda":  0},
+      {"type": "potential", "lambda": 0},
+      {"type": "flux_preservation", "lambda":  0}
+    ],
+    "check_val_every_n_epoch": 10
+  }
+}
\ No newline at end of file
diff --git a/config_yaml/multi_height/muram_mfr.yaml b/config_yaml/multi_height/muram_mfr.yaml
new file mode 100644
index 0000000..305a2fc
--- /dev/null
+++ b/config_yaml/multi_height/muram_mfr.yaml
@@ -0,0 +1,53 @@
+---
+base_path: "/glade/work/rjarolim/nf2/muram_mfr/muram_mfr_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_v01"
+logging:
+  project: muram_mfr
+  name: muram MFR
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_1.000.474000"
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.010.474000"
+      height_mapping: {z: 1.0, z_min: 0, z_max: 2}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000100.474000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 4}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: {z: 6.0, z_min: 0, z_max: 12}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02, boundary_03, boundary_04]
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [boundary_02, boundary_03, boundary_04]
+  coordinate_transform:
+      type: height
+      ds_id: [boundary_02, boundary_03, boundary_04]
+      validation_ds_id: [validation_boundary_02, validation_boundary_03, validation_boundary_04]
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/multi_height/muram_mfr_disambiguation.yaml b/config_yaml/multi_height/muram_mfr_disambiguation.yaml
new file mode 100644
index 0000000..8600bad
--- /dev/null
+++ b/config_yaml/multi_height/muram_mfr_disambiguation.yaml
@@ -0,0 +1,58 @@
+---
+base_path: "/glade/work/rjarolim/nf2/muram_mfr/muram_mfr_disambiguation_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_disambiguation_v02"
+logging:
+  project: muram_mfr
+  name: muram MFR disambiguation
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_1.000.474000"
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.010.474000"
+      height_mapping: {z: 1.0, z_min: 0, z_max: 2}
+      los_trv_azi_transform: True
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000100.474000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 4}
+      los_trv_azi_transform: True
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: {z: 6.0, z_min: 0, z_max: 12}
+      los_trv_azi_transform: True
+#      batch_size: 2048
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0
+      ds_id: boundary_01
+    - type: boundary_los_trv_azi
+      name: boundary_01
+      lambda: 1.0
+      ds_id: [boundary_02, boundary_03, boundary_04]
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda:  1.0e-1
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [boundary_02, boundary_03, boundary_04]
+  coordinate_transform:
+      type: height
+      ds_id: [boundary_02, boundary_03, boundary_04]
+      validation_ds_id: [validation_boundary_02, validation_boundary_03, validation_boundary_04]
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/multi_height/muram_mfr_single_height.yaml b/config_yaml/multi_height/muram_mfr_single_height.yaml
new file mode 100644
index 0000000..eba0d9c
--- /dev/null
+++ b/config_yaml/multi_height/muram_mfr_single_height.yaml
@@ -0,0 +1,32 @@
+---
+base_path: "/glade/work/rjarolim/nf2/muram_mfr/muram_mfr_single_height_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_single_height_v01"
+logging:
+  project: muram_mfr
+  name: muram MFR single-height
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_1.000.474000"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0
+      ds_id: boundary_01
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/multi_height/muram_sma.yaml b/config_yaml/multi_height/muram_sma.yaml
new file mode 100644
index 0000000..3cdd9d3
--- /dev/null
+++ b/config_yaml/multi_height/muram_sma.yaml
@@ -0,0 +1,54 @@
+---
+base_path: "/glade/work/rjarolim/nf2/muram_mfr/muram_sma_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_sma_v01"
+logging:
+  project: muram_mfr
+  name: muram SMA
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_1.000.399000"
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_0.010.399000"
+      height_mapping: {z: 1.0, z_min: 0, z_max: 2}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_0.000100.399000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 4}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_0.000001.399000"
+      height_mapping: {z: 6.0, z_min: 0, z_max: 12}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02, boundary_03, boundary_04]
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [boundary_02, boundary_03, boundary_04]
+  coordinate_transform:
+      type: height
+      ds_id: [boundary_02, boundary_03, boundary_04]
+      validation_ds_id: [validation_boundary_02, validation_boundary_03, validation_boundary_04]
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/multi_height/muram_sma_single_height.yaml b/config_yaml/multi_height/muram_sma_single_height.yaml
new file mode 100644
index 0000000..77903c2
--- /dev/null
+++ b/config_yaml/multi_height/muram_sma_single_height.yaml
@@ -0,0 +1,32 @@
+---
+base_path: "/glade/work/rjarolim/nf2/muram_mfr/muram_sma_single_height_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_sma_single_height_v01"
+logging:
+  project: muram_mfr
+  name: muram SMA single-height
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_1.000.399000"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0
+      ds_id: boundary_01
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/multi_height/muram_xflare.yaml b/config_yaml/multi_height/muram_xflare.yaml
new file mode 100644
index 0000000..ccd76f8
--- /dev/null
+++ b/config_yaml/multi_height/muram_xflare.yaml
@@ -0,0 +1,53 @@
+---
+base_path: "/glade/work/rjarolim/nf2/multi_height/muram_xflare_v08"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_xflare_v08"
+logging:
+  project: multi-height
+  name: muram X-flare
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/ASD_2023/Run_Corona_1728x1024x1024_ASD_open/2D/tau_slice_1.000.1300000"
+      height_mapping: {z: 0.5, z_min: 0, z_max: 1}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/ASD_2023/Run_Corona_1728x1024x1024_ASD_open/2D/tau_slice_0.001.1300000"
+      height_mapping: {z: 1.0, z_min: 0, z_max: 2}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/ASD_2023/Run_Corona_1728x1024x1024_ASD_open/2D/tau_slice_0.000100.1300000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 4}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/ASD_2023/Run_Corona_1728x1024x1024_ASD_open/2D/tau_slice_0.000001.1300000"
+      height_mapping: {z: 6.0, z_min: 0, z_max: 12}
+      los_trv_azi_transform: False
+#      batch_size: 2048
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 10
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+2, end: 1.0, iterations: 1.0e+5}
+      ds_id: [boundary_01, boundary_02, boundary_03, boundary_04]
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0e-1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: min_height
+      name: height
+      lambda: {start: 1.0e-2, end: 1.0e-6, iterations: 1.0e+5}
+      ds_id: [boundary_01, boundary_02, boundary_03, boundary_04]
+  coordinate_transform:
+      type: height
+      ds_id: [boundary_01, boundary_02, boundary_03, boundary_04]
+      validation_ds_id: [validation_boundary_01, validation_boundary_02, validation_boundary_03, validation_boundary_04]
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/sharp/13229.yaml b/config_yaml/sharp/13229.yaml
new file mode 100644
index 0000000..a306602
--- /dev/null
+++ b/config_yaml/sharp/13229.yaml
@@ -0,0 +1,32 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13229_custom_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13229_v01"
+logging:
+  project: nf2-cartesian
+  name: NOAA 13229
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13229_custom/hmi.Bptr_720s_mod.20230224_140000_TAI.JSOC_20240609_003437_IN.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+  max_height: 150
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664.yaml b/config_yaml/sharp/13664.yaml
new file mode 100644
index 0000000..ccad7c1
--- /dev/null
+++ b/config_yaml/sharp/13664.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/init"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13664_v01"
+logging:
+  project: nf2-cartesian
+  name: NOAA 13664
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/hmi.Bptr_720s_mod.20240504_120000_TAI.JSOC_20240521_001951_IN.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_ensemble/run01.yaml b/config_yaml/sharp/13664_ensemble/run01.yaml
new file mode 100644
index 0000000..720ca58
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/run01.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/ensemble/run01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/ensemble/run01"
+logging:
+  project: Ensemble 13664
+  name: run01
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_ensemble/run02.yaml b/config_yaml/sharp/13664_ensemble/run02.yaml
new file mode 100644
index 0000000..511c1d9
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/run02.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/ensemble/run02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/ensemble/run02"
+logging:
+  project: Ensemble 13664
+  name: run02
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_ensemble/run03.yaml b/config_yaml/sharp/13664_ensemble/run03.yaml
new file mode 100644
index 0000000..8c5f6e1
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/run03.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/ensemble/run03"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/ensemble/run03"
+logging:
+  project: Ensemble 13664
+  name: run03
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_ensemble/run04.yaml b/config_yaml/sharp/13664_ensemble/run04.yaml
new file mode 100644
index 0000000..36f4de5
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/run04.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/ensemble/run04"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/ensemble/run04"
+logging:
+  project: Ensemble 13664
+  name: run04
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_ensemble/run05.yaml b/config_yaml/sharp/13664_ensemble/run05.yaml
new file mode 100644
index 0000000..d918e2f
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/run05.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/ensemble/run05"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/ensemble/run05"
+logging:
+  project: Ensemble 13664
+  name: run05
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/20240509_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13664_series.yaml b/config_yaml/sharp/13664_series.yaml
new file mode 100644
index 0000000..b34b943
--- /dev/null
+++ b/config_yaml/sharp/13664_series.yaml
@@ -0,0 +1,24 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664/series"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13664_series_v01"
+meta_path: "/glade/work/rjarolim/nf2/sharp/13664/init/last.ckpt"
+logging:
+  project: AR 13664
+  name: series
+data:
+  type: sharp
+  data_path: "/glade/work/rjarolim/data/nf2/13664"
+  num_workers: 8
+  iterations: 2.0e+3
+model:
+  type: vector_potential
+  dim: 256
+training:
+  check_val_every_n_epoch: 5 # validation plots in 1h steps
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
\ No newline at end of file
diff --git a/config_yaml/sharp/377.yaml b/config_yaml/sharp/377.yaml
new file mode 100644
index 0000000..ebae146
--- /dev/null
+++ b/config_yaml/sharp/377.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/377_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_377_v02"
+logging:
+  project: nf2-cartesian
+  name: vector potential
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp_err.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0 #{start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: force_free
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/377_series.yaml b/config_yaml/sharp/377_series.yaml
new file mode 100644
index 0000000..6389dd1
--- /dev/null
+++ b/config_yaml/sharp/377_series.yaml
@@ -0,0 +1,26 @@
+---
+base_path: "/glade/work/rjarolim/nf2/cartesian_series/sharp_377_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_series_v01"
+meta_path: "/glade/work/rjarolim/nf2/cartesian_series/sharp_377_v01/last.ckpt"
+logging:
+  project: nf2-cartesian
+  name: series
+data:
+  type: sharp
+  data_path: "/glade/work/rjarolim/data/nf2/377"
+  num_workers: 8
+  iterations: 2.0e+3
+model:
+  type: vector_potential
+  dim: 256
+  activation: swish
+training:
+  epochs: 1.0e+3
+  check_val_every_n_epoch: 5 # validation plots in 1h steps
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: [boundary, potential]
+    - type: force_free
+      lambda: 1.0e-1
\ No newline at end of file
diff --git a/config_yaml/spherical/2154.yaml b/config_yaml/spherical/2154.yaml
new file mode 100644
index 0000000..d778b34
--- /dev/null
+++ b/config_yaml/spherical/2154.yaml
@@ -0,0 +1,82 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2154_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2154_v01"
+logging:
+  project: spherical
+  name: "2154"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 2.087, 4.915 ]
+        latitude_range: [  0.186, 2.699 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 2.087, 4.915 ]
+      latitude_range: [  0.186, 2.699 ]
+  validation_configs:
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_120000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 2.087, 4.915 ]
+        latitude_range: [  0.186, 2.699 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 2.087, 4.915 ]
+      latitude_range: [  0.186, 2.699 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 2.087, 4.915 ]
+      latitude_range: [  0.186, 2.699 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: full_disk
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/2173.yaml b/config_yaml/spherical/2173.yaml
new file mode 100644
index 0000000..dabc768
--- /dev/null
+++ b/config_yaml/spherical/2173.yaml
@@ -0,0 +1,82 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_v03"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_v03"
+logging:
+  project: spherical
+  name: "2173"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.446, 4.273 ]
+        latitude_range: [  0.424, 2.938 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+  validation_configs:
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.446, 4.273 ]
+        latitude_range: [  0.424, 2.938 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1 #{start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: full_disk
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
\ No newline at end of file
diff --git a/config_yaml/spherical/2173_full.yaml b/config_yaml/spherical/2173_full.yaml
new file mode 100644
index 0000000..a3bbfb0
--- /dev/null
+++ b/config_yaml/spherical/2173_full.yaml
@@ -0,0 +1,85 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_full_divergence"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_full_divergence"
+logging:
+  project: spherical
+  name: "2173 full"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+#        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+#        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+#        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+    - type: map
+      ds_id: synoptic
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/hmi.b_synoptic.2173.Bp.fits"
+      mask_configs:
+        type: reference
+        file: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 16384
+  validation_configs:
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.446, 4.273 ]
+        latitude_range: [  0.424, 2.938 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1.0e+1 #{start: 1.0e+3, end: 10.0, iterations: 5.0e+4}
+      ds_id: [full_disk, synoptic]
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/2173_subframe.yaml b/config_yaml/spherical/2173_subframe.yaml
new file mode 100644
index 0000000..8fbc771
--- /dev/null
+++ b/config_yaml/spherical/2173_subframe.yaml
@@ -0,0 +1,82 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_subframe_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_subframe_v01"
+logging:
+  project: spherical
+  name: "2173 subframe"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 2.336, 3.384 ]
+        latitude_range: [  1.157, 2.205 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 2.336, 3.384 ]
+      latitude_range: [  1.157, 2.205 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_170000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 2.336, 3.384 ]
+        latitude_range: [  1.157, 2.205 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 2.336, 3.384 ]
+      latitude_range: [  1.157, 2.205 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 2.336, 3.384 ]
+      latitude_range: [  1.157, 2.205 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: AR
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/2267.yaml b/config_yaml/spherical/2267.yaml
new file mode 100644
index 0000000..828b122
--- /dev/null
+++ b/config_yaml/spherical/2267.yaml
@@ -0,0 +1,73 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2267_v04"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2267_v04"
+logging:
+  project: spherical
+  name: "2267 radial"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp_err.fits"
+    - type: map
+      ds_id: synoptic
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/hmi.synoptic_mr_polfil_720s.2267.Mr_polfil.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/hmi.b_synoptic.2267.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/hmi.b_synoptic.2267.Bp.fits"
+      mask_configs:
+        type: reference
+        file: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br.fits"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 16384
+      name: random
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp_err.fits"
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+    - type: sphere
+      ds_id: sphere
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 65536
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: [synoptic, AR]
+    - type: force_free
+      lambda: 0.1
+    - type: radial
+      lambda: 1.0e-4
+      ds_id: [random]
+      base_radius: 7.246875
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/spherical/2267_post_eruption.yaml b/config_yaml/spherical/2267_post_eruption.yaml
new file mode 100644
index 0000000..70cb8c8
--- /dev/null
+++ b/config_yaml/spherical/2267_post_eruption.yaml
@@ -0,0 +1,73 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2267_post_eruption_v03"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2267_post_eruption_v03"
+logging:
+  project: spherical
+  name: "2267 post eruption"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bp_err.fits"
+    - type: map
+      ds_id: synoptic
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/hmi.synoptic_mr_polfil_720s.2267.Mr_polfil.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/hmi.b_synoptic.2267.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/hmi.b_synoptic.2267.Bp.fits"
+      mask_configs:
+        type: reference
+        file: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Br.fits"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 16384
+      name: random
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230225_000000_TAI.Bp_err.fits"
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+    - type: sphere
+      ds_id: sphere
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 65536
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: [synoptic, AR]
+    - type: force_free
+      lambda: 0.1
+    - type: radial
+      lambda: 1.0e-4
+      ds_id: [random]
+      base_radius: 7.246875
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/spherical/2267_subframe.yaml b/config_yaml/spherical/2267_subframe.yaml
new file mode 100644
index 0000000..3aba871
--- /dev/null
+++ b/config_yaml/spherical/2267_subframe.yaml
@@ -0,0 +1,79 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2267_subframe_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2267_subframe_v02"
+logging:
+  project: spherical
+  name: "2267 subframe"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.0, 8.0 ]
+        latitude_range: [ 0, 1.57 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 5.0, 8.0 ]
+      latitude_range: [ 0, 1.57 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2267/full_disk/hmi.b_720s.20230224_190000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.0, 8.0 ]
+        latitude_range: [ 0, 1.57 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 5.0, 8.0 ]
+      latitude_range: [ 0, 1.57 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 5.0, 8.0 ]
+      latitude_range: [ 0, 1.57 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: AR
+    - type: force_free
+      lambda: 0.1
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+  check_val_every_n_epoch: 1
diff --git a/config_yaml/spherical/2283_00.yaml b/config_yaml/spherical/2283_00.yaml
new file mode 100644
index 0000000..819d6bb
--- /dev/null
+++ b/config_yaml/spherical/2283_00.yaml
@@ -0,0 +1,78 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2283_00_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2283_00_v01"
+logging:
+  project: spherical
+  name: '2283'
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.673, 4.501 ]
+        latitude_range: [ 0.398, 2.911 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 1.673, 4.501 ]
+      latitude_range: [ 0.398, 2.911 ]
+  validation_configs:
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.673, 4.501 ]
+        latitude_range: [ 0.398, 2.911 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 1.673, 4.501 ]
+      latitude_range: [ 0.398, 2.911 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: full_disk
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/2283_04.yaml b/config_yaml/spherical/2283_04.yaml
new file mode 100644
index 0000000..7a67eee
--- /dev/null
+++ b/config_yaml/spherical/2283_04.yaml
@@ -0,0 +1,78 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2283_04_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2283_04_v01"
+logging:
+  project: spherical
+  name: '2283'
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.673, 4.501 ]
+        latitude_range: [ 0.398, 2.911 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 1.673, 4.501 ]
+      latitude_range: [ 0.398, 2.911 ]
+  validation_configs:
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_040000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.673, 4.501 ]
+        latitude_range: [ 0.398, 2.911 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 1.673, 4.501 ]
+      latitude_range: [ 0.398, 2.911 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: full_disk
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/377_embedded.yaml b/config_yaml/spherical/377_embedded.yaml
new file mode 100644
index 0000000..13b5d5d
--- /dev/null
+++ b/config_yaml/spherical/377_embedded.yaml
@@ -0,0 +1,97 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/377_embedded_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_embedded_v01"
+logging:
+  project: spherical
+  name: '377 embedded'
+data:
+  type: spherical
+  train_configs:
+    # full disk configuration --> lower sampling
+    - type: map
+      ds_id: full_disk
+      batch_size: 2048
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+    - type: random_spherical
+      ds_id: random_full_disk
+      radius_range: [1.0, 1.3]
+      batch_size: 2048
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random_AR
+      batch_size: 8192
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+  max_radius: 1.14
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: [AR, full_disk]
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: [ random_full_disk, random_AR ]
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: [ random_full_disk, random_AR ]
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/377_potential.yaml b/config_yaml/spherical/377_potential.yaml
new file mode 100644
index 0000000..21d90af
--- /dev/null
+++ b/config_yaml/spherical/377_potential.yaml
@@ -0,0 +1,111 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/377_potential_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_potential_v01"
+logging:
+  project: spherical
+  name: '377 potential'
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.263, 8.090 ]
+        latitude_range: [ 0.434, 2.947 ]
+        invert: true
+    - type: pfss_boundary
+      ds_id: pfss_boundary
+      Br: "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits"
+      insert: ["/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"]
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      mask:
+      - longitude_range: [ 5.263, 8.090 ]
+      - latitude_range: [ 0.434, 2.947 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.263, 8.090 ]
+        latitude_range: [ 0.434, 2.947 ]
+        invert: true
+    - type: pfss_boundary
+      ds_id: pfss_boundary_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      Br: "/glade/work/rjarolim/data/global/fd_2106/hmi.synoptic_mr_polfil_720s.2106.Mr_polfil.fits"
+      insert: [ "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits" ]
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: full_disk
+    - type: boundary
+      name: pfss_boundary
+      lambda: 0.1 #1
+      ds_id: pfss_boundary
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/377_spherical.yaml b/config_yaml/spherical/377_spherical.yaml
new file mode 100644
index 0000000..b0ea1b9
--- /dev/null
+++ b/config_yaml/spherical/377_spherical.yaml
@@ -0,0 +1,99 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/377_v05"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_v05"
+logging:
+  project: spherical
+  name: '377 spherical'
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.263, 8.090 ]
+        latitude_range: [ 0.434, 2.947 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: map
+      ds_id: full_disk_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 5.263, 8.090 ]
+        latitude_range: [ 0.434, 2.947 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 5.263, 8.090 ]
+      latitude_range: [ 0.434, 2.947 ]
+  max_radius: 1.3
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
+      ds_id: full_disk
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/spherical/377_subframe.yaml b/config_yaml/spherical/377_subframe.yaml
new file mode 100644
index 0000000..09df26b
--- /dev/null
+++ b/config_yaml/spherical/377_subframe.yaml
@@ -0,0 +1,82 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/377_subframe_v04"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_subframe_v04"
+logging:
+  project: spherical
+  name: '377 subframe'
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: AR
+      batch_size: 4096
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+  validation_configs:
+    - type: map
+      ds_id: AR_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2106/full_disk/hmi.b_720s.20110215_000000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 0.3, 0.9 ]
+        latitude_range: [ 1.74, 2.14 ]
+        invert: true
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 0.3, 0.9 ]
+      latitude_range: [ 1.74, 2.14 ]
+  max_radius: 1.14
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 1000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: AR
+    - type: energy_gradient
+      lambda: 1.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
diff --git a/config_yaml/sst/7310.yaml b/config_yaml/sst/7310.yaml
new file mode 100644
index 0000000..73111d6
--- /dev/null
+++ b/config_yaml/sst/7310.yaml
@@ -0,0 +1,70 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/7310_v04"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/7310_v04"
+logging:
+  project: nf2-sst
+  name: "7310 SST"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/converted/BLOS_2018-09-30T09:22:00_FeI6173_v3_Robert.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/converted/BTRV_2018-09-30T09:22:00_FeI6173_v3_Robert.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/converted/BAZI_2018-09-30T09:22:00_FeI6173_v3_Robert.fits"
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/converted/BLOS_2018-09-30T09:22:00_NaI5896_v3_Robert.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/converted/BTRV_2018-09-30T09:22:00_NaI5896_v3_Robert.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/converted/BAZI_2018-09-30T09:22:00_NaI5896_v3_Robert.fits"
+      height_mapping: {z: 1.0, z_min: 0.0, z_max: 5.0}
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/converted/BLOS_2018-09-30T09:22:00_CaII8542_v3_Robert.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/converted/BTRV_2018-09-30T09:22:00_CaII8542_v3_Robert.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/converted/BAZI_2018-09-30T09:22:00_CaII8542_v3_Robert.fits"
+      height_mapping: {z: 2.0, z_min: 0.0, z_max: 10.0}
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/converted/hmi_los_masked.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/converted/hmi_trv_masked.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/converted/hmi_azi_masked.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02, boundary_03]
+    - type: boundary_los_trv_azi
+      name: boundary_hmi
+      lambda: 0.1
+      ds_id: boundary_04
+      disambiguate: false
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: divergence
+      lambda: 1.0e-2
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [boundary_02, boundary_03]
+  coordinate_transform:
+      type: height
+      ds_id: [boundary_02, boundary_03]
+      validation_ds_id: [validation_boundary_02, validation_boundary_03]
+  check_val_every_n_epoch: 1
\ No newline at end of file
diff --git a/config_yaml/sst/7310_hmi.yaml b/config_yaml/sst/7310_hmi.yaml
new file mode 100644
index 0000000..0c48ec3
--- /dev/null
+++ b/config_yaml/sst/7310_hmi.yaml
@@ -0,0 +1,38 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/7310_hmi_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/7310_hmi_v02"
+logging:
+  project: nf2-sst
+  name: "7310 HMI VP"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/converted/hmi_los.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/converted/hmi_trv.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/converted/hmi_azi.fits"
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+  activation: sine
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary_hmi
+      lambda: 1.0
+      ds_id: boundary_01
+      disambiguate: false
+    - type: boundary
+      name: potential_boundary
+      lambda: 1.0
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: divergence
+      lambda: 1.0e-2
+  check_val_every_n_epoch: 1
\ No newline at end of file
diff --git a/example/extrapolation.ipynb b/example/extrapolation.ipynb
deleted file mode 100644
index 8afe9b0..0000000
--- a/example/extrapolation.ipynb
+++ /dev/null
@@ -1,1179 +0,0 @@
-{
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "NY2f9O3T_Ovk"
-      },
-      "source": [
-        "# Neural Network Force-Free magnetic field extrapolation - NF2\n",
-        "<img src=\"https://github.com/RobertJaro/NF2/blob/main/images/logo.jpg?raw=true\" width=\"150\" height=\"150\">\n",
-        "\n",
-        "This notebook uses NF2 to download and extrapolate vector magnetograms from SDO/HMI. Active regions and dates can be selected below.\n",
-        "\n",
-        "We use NF2 to extrapolate a time-series and to analyze the resulting magnetic energy build-up and release processes (e.g., flares).\n",
-        "\n",
-        "GitHub Page: https://github.com/RobertJaro/NF2\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1jJgVkX-MXsr"
-      },
-      "source": [
-        "**Use**: A python notebook consists of code cells and embedded text descriptions. If you are using Google Colab all computations and installations are performed online, using free resources. You can execute the notebook cell-by-cell (hit `shift`+`return`), or run the full notebook by selecting `Runtime` --> `Run all` from the menu. Use the forms below to select the active region, date, and training parameters. For regular use you can keep the default training parameters."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "vGdTQSN6GA6z"
-      },
-      "source": [
-        "## Installation and Imports"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "deEBHLxM0Uqb"
-      },
-      "source": [
-        "NF2 can be directly installed from the GitHub repository."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "l5_9wsH1cmsh"
-      },
-      "outputs": [],
-      "source": [
-        "!pip install git+https://github.com/RobertJaro/NF2.git"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "X1fOxaJS5xH6"
-      },
-      "outputs": [],
-      "source": [
-        "# general imports\n",
-        "import glob\n",
-        "import os\n",
-        "from datetime import datetime, timedelta\n",
-        "from dateutil.parser import parse\n",
-        "import shutil\n",
-        "import gdown\n",
-        "\n",
-        "# download\n",
-        "import drms\n",
-        "from urllib import request\n",
-        "\n",
-        "# data processing\n",
-        "import numpy as np\n",
-        "from sunpy.map import Map\n",
-        "from sunpy.net import Fido\n",
-        "from sunpy.net import attrs as a\n",
-        "\n",
-        "# deep learning\n",
-        "import torch\n",
-        "from pytorch_lightning import Trainer\n",
-        "from pytorch_lightning.loggers import WandbLogger\n",
-        "from pytorch_lightning.callbacks import ModelCheckpoint, LambdaCallback\n",
-        "import wandb\n",
-        "\n",
-        "# NF2\n",
-        "from nf2.train.module import NF2Module, save\n",
-        "from nf2.train.data_loader import SHARPDataModule, SHARPSeriesDataModule\n",
-        "from nf2.data.download import download_HARP_series, find_HARP, download_euv\n",
-        "from nf2.evaluation.unpack import load_cube\n",
-        "from nf2.evaluation.metric import *\n",
-        "from nf2.evaluation.energy import get_free_mag_energy\n",
-        "from nf2.evaluation.series import evaluate_nf2_series\n",
-        "from nf2.evaluation.flares import _calculate_free_energy, get_integrated_euv_map, load_B_map\n",
-        "\n",
-        "# visualization\n",
-        "from matplotlib import pyplot as plt"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dGrZwKgwGEro"
-      },
-      "source": [
-        "## Data download"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "oZ8WH-7ABW8J"
-      },
-      "source": [
-        "We start by downloading a series of SDO/HMI vector magnetograms. Use the forms below to set your email, and specify the active region that you are interested in (SHARP + date)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Ufn-xIQBCKS8"
-      },
-      "source": [
-        "Downloading data requires an active registration at JSOC. http://jsoc.stanford.edu/ajax/register_email.html"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "wJRSB1NV--Bb"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Download Credentials\n",
-        "jsoc_email = 'robert.jarolim@uni-graz.at' #@param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "92DkBePsCiTd"
-      },
-      "outputs": [],
-      "source": [
-        "client = drms.Client(email=jsoc_email, verbose=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "JzBnLWbEB88F"
-      },
-      "source": [
-        "If you are looking for a NOAA active region, you can use the function below to search for the corrsponding HARP number."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2zVLpXBf2oQN",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Start date of observation series\n",
-        "year = 2017 #@param {type:\"integer\"}\n",
-        "month = 9 #@param {type:\"integer\"}\n",
-        "day = 6 #@param {type:\"integer\"}\n",
-        "hour = 8 #@param {type:\"integer\"}\n",
-        "minute = 36 #@param {type:\"number\"}\n",
-        "\n",
-        "date = datetime(year, month, day, hour, minute)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "OOb-ZnK40qt-",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Find SHARP number for given NOAA\n",
-        "noaa_nums = [12673] #@param {type:\"\"}\n",
-        "\n",
-        "find_HARP(date, noaa_nums, client)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "nDUBoZ3-86en",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Active Region\n",
-        "#@markdown (add the number from above here)\n",
-        "sharp_nr = 7115 #@param {type:\"number\"}\n",
-        "duration = '3h' # @param {type:\"string\"}\n",
-        "\n",
-        "date = datetime(year, month, day, hour, minute)\n",
-        "\n",
-        "download_dir = 'AR_7115' #@param {type:\"string\"}\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "p5FqGAP55UUO"
-      },
-      "source": [
-        "Start download for time-series."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "1W22DbZHzYm4"
-      },
-      "outputs": [],
-      "source": [
-        "download_HARP_series(sharp_nr, date, duration, download_dir, client)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Let's check the first map in the series."
-      ],
-      "metadata": {
-        "id": "fFQE-dKj5w6p"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "Map(glob.glob(os.path.join(download_dir, '*Br.fits'))[0]).peek()"
-      ],
-      "metadata": {
-        "id": "eV19mXCB5vd9"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "SlfpqlAyccTG"
-      },
-      "source": [
-        "## Set parameters"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ovL6D7Az7_CO",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Paths\n",
-        "base_path = None #@param {type:\"string\"}\n",
-        "base_path = 'ar_%d_%s' % (sharp_nr, date.isoformat('T')) if base_path is None else base_path\n",
-        "data_path = None #@param {type:\"string\"}\n",
-        "data_path = download_dir if data_path is None else data_path\n",
-        "\n",
-        "os.makedirs(base_path, exist_ok=True)\n",
-        "\n",
-        "save_path = os.path.join(base_path, 'extrapolation_result.nf2')"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ZLeXpnuP6DHq",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Data Parameters\n",
-        "bin = 2 #@param {type:\"integer\"}\n",
-        "spatial_norm = 160  #@param {type:\"integer\"}\n",
-        "height = 160  #@param {type:\"integer\"}\n",
-        "b_norm = 2500  #@param {type:\"number\"}\n",
-        "d_slice =  None#@param {type:\"raw\"}\n",
-        "potential_boundary = True #@param {type:\"boolean\"}\n",
-        "\n",
-        "\n",
-        "boundary = {'type': 'potential', 'strides': 4} if potential_boundary else {'type': 'open'}\n",
-        "\n",
-        "data_args = {\"data_path\": data_path,\n",
-        "             \"height\":height,\n",
-        "             \"spatial_norm\": spatial_norm,\n",
-        "             \"b_norm\": b_norm,\n",
-        "             \"work_directory\": base_path,\n",
-        "             \"bin\": bin,\n",
-        "             \"Mm_per_pixel\": 0.72,\n",
-        "             \"slice\": d_slice,\n",
-        "             \"boundar\": boundary\n",
-        "             }"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ZPAcHR4K6sU8",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Model Parameters\n",
-        "#@markdown `dim = 64` &rarr; tiny\n",
-        "\n",
-        "#@markdown `dim = 256` &rarr; regular\n",
-        "\n",
-        "dim = 256 #@param {type:\"number\"}\n",
-        "vector_potential = False #@param {type:\"boolean\"}\n",
-        "\n",
-        "model_args = {\"dim\": dim, \"use_vector_potential\": vector_potential}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "4KoeQGTp6t1B",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Training Parameters\n",
-        "lambda_div = 0.1 #@param {type:\"number\"}\n",
-        "lambda_ff = 0.1 #@param {type:\"number\"}\n",
-        "\n",
-        "iterations = 8e4 #@param {type:\"number\"}\n",
-        "iterations = int(iterations)\n",
-        "\n",
-        "validation_interval = 1e4 #@param {type:\"number\"}\n",
-        "validation_interval = int(validation_interval)\n",
-        "\n",
-        "\n",
-        "batch_size = 1e4 #@param {type:\"number\"}\n",
-        "batch_size = int(batch_size)\n",
-        "\n",
-        "data_args['iterations'] = iterations\n",
-        "data_args['batch_size'] = batch_size\n",
-        "training_args = {\"lambda_div\": lambda_div,\n",
-        "              \"lambda_ff\": lambda_ff,}\n",
-        "# combine args\n",
-        "config = {'data': data_args, 'model': model_args, 'training': training_args}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8KF8grbpIlnU"
-      },
-      "source": [
-        "## Training - Extrapolation"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Initialize the Weights and Biases (wandb) logger. The training results will be automatically uploaded during training and model checkpoints can be used to continue the model training."
-      ],
-      "metadata": {
-        "id": "4BhHTCe7N3iM"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "BQPv88a17_3Z"
-      },
-      "outputs": [],
-      "source": [
-        "wandb_logger = WandbLogger(project='nf2', name=str(sharp_nr), dir=base_path, log_model=\"all\")\n",
-        "wandb_logger.experiment.config.update(config, allow_val_change=True)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "#@title Continue training from WandB\n",
-        "#@markdown Enter your credentials if you want to continue an inturrpted run. This will restore the latest model and continue the training.\n",
-        "wandb_id = \"\" #@param {type:\"string\"}\n",
-        "wandb_entity = \"\" #@param {type:\"string\"}\n",
-        "wandb_project = \"\" #@param {type:\"string\"}\n",
-        "if wandb_id != \"\" and wandb_id is not None:\n",
-        "    checkpoint_reference = f\"{wandb_entity}/{wandb_project}/model-{wandb_id}:latest\"\n",
-        "    artifact = wandb_logger.use_artifact(checkpoint_reference, artifact_type=\"model\")\n",
-        "    artifact.download(root=base_path)\n",
-        "    shutil.move(os.path.join(base_path, 'model.ckpt'), os.path.join(base_path, 'last.ckpt'))\n",
-        "    data_args['plot_overview'] = False  # skip overview plot for restored model"
-      ],
-      "metadata": {
-        "cellView": "form",
-        "id": "Avh_MFrzdLVr"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "QXZj5mZY8Aqu"
-      },
-      "source": [
-        "For logging we use WandB which allows us to monitor the traing process and saves our results."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "WD8cJ9ZEC7E2"
-      },
-      "source": [
-        "Before we start the extrapolation we need to process the downloaded data. The model requires the vector magnetogram and the corresponding error map. For this we use a data loader that automatically handles the processing of SHARPS according to the parameters set above."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "j6Y-v3dA7GSc"
-      },
-      "outputs": [],
-      "source": [
-        "data_module = SHARPDataModule(**data_args)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sc7PUWtJV7qk"
-      },
-      "source": [
-        "We initialize the NF2 model with the parameters from above."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ZfEibAPOIoQS"
-      },
-      "outputs": [],
-      "source": [
-        "validation_settings = {'cube_shape': data_module.cube_dataset.coords_shape,\n",
-        "                       'gauss_per_dB': b_norm,\n",
-        "                       'Mm_per_ds': data_module.Mm_per_pixel * spatial_norm}\n",
-        "\n",
-        "nf2 = NF2Module(validation_settings, **model_args, **training_args)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "GdkZvLIuWMqB"
-      },
-      "source": [
-        "During training we log model states with callbacks that are used after every validation step. The models are then automatically synched to wandb. The extrapolation result (`.nf2` file) can be found in the file explorer on the left."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "sP_20sNcV30S"
-      },
-      "outputs": [],
-      "source": [
-        "save_callback = LambdaCallback(\n",
-        "    on_validation_end=lambda *args: save(save_path, nf2.model, data_module, config, nf2.height_mapping_model))\n",
-        "checkpoint_callback = ModelCheckpoint(dirpath=base_path,\n",
-        "                                      every_n_train_steps=validation_interval,\n",
-        "                                      save_last=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "CwAd1rvJHsJJ"
-      },
-      "source": [
-        "We initiliaze the traininer from pytorch lightning which handles the individual components."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "f_yZENeDXh60"
-      },
-      "outputs": [],
-      "source": [
-        "n_gpus = torch.cuda.device_count()\n",
-        "trainer = Trainer(max_epochs=1,\n",
-        "                  logger=wandb_logger,\n",
-        "                  devices=n_gpus if n_gpus >= 1 else None,\n",
-        "                  accelerator='gpu' if n_gpus >= 1 else None,\n",
-        "                  strategy='dp' if n_gpus > 1 else None,\n",
-        "                  num_sanity_val_steps=0,\n",
-        "                  val_check_interval=validation_interval,\n",
-        "                  gradient_clip_val=0.1,\n",
-        "                  callbacks=[checkpoint_callback, save_callback], )"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "rCKi-HwKZY3h"
-      },
-      "source": [
-        "Now we can start our model training. The training progress is automatically logged to wandb. Follow the link above to see the status and intermediate results."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "rs2CgBFmYY0e"
-      },
-      "outputs": [],
-      "source": [
-        "trainer.fit(nf2, data_module, ckpt_path='last')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "As last step we save the trained model. If you continue an interrupted training run, excecude the code below to restore the nf2 file."
-      ],
-      "metadata": {
-        "id": "8-BSQlguOjF1"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "save(save_path, nf2.model, data_module, config, nf2.height_mapping_model)"
-      ],
-      "metadata": {
-        "id": "ICNqfLUOOiXQ"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "YT5wsu5xa5Nz"
-      },
-      "source": [
-        "## Evaluation"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "U-YDUsO5Vq-L"
-      },
-      "source": [
-        "If you want to skip the training you can use the line below to download an example solution."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "wmjG-hDua2wc",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Download sample extrapolation result\n",
-        "download = False #@param {type:\"boolean\"}\n",
-        "tiny = False #@param {type:\"boolean\"}\n",
-        "if download:\n",
-        "  if tiny:\n",
-        "    gdown.download('https://drive.google.com/uc?id=1BwvEtsw10pE18RDfSPnCKy3476FqV47_', save_path)\n",
-        "  else:\n",
-        "    gdown.download('https://drive.google.com/uc?id=1m57uDWSOiTd0cYUpZGr7nQ5HO00iTePl', save_path)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ZklkaZOjHJUg"
-      },
-      "source": [
-        "It is important that the resutls are checked after the extrapolation is finished. The NF2 package provides metrics that can be used to estimate the quality. Both the divergence and the angle between the currents and magnetic field should be small. (ideally mean divergence `<0.1` and sigma `<10` degree; the exact values depend on the selected active region and lambda weights)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "-EEaBAvoIwHT"
-      },
-      "source": [
-        "From the save file we can load the full simulation volume."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "KzP-QmwWHiHC"
-      },
-      "outputs": [],
-      "source": [
-        "b = load_cube(save_path, progress=True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Q3I5u8oSSbKA"
-      },
-      "source": [
-        "The loaded mesh can be used to compute the divergence and the sigma angle."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "Pnj8vZ1YSarz"
-      },
-      "outputs": [],
-      "source": [
-        "n_div = (np.abs(divergence(b)) / vector_norm(b)).mean()\n",
-        "theta = weighted_theta(b)\n",
-        "\n",
-        "print('DIVERGENCE [1/pix]: %.04f; THETA [deg] %.04f' % (n_div, theta))"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "VtBgyVytQW7Q"
-      },
-      "source": [
-        "We can also compute maps of integrated current density."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "YZYuXgvPQcqX"
-      },
-      "outputs": [],
-      "source": [
-        "j = curl(b)\n",
-        "plt.figure(figsize=(10,4))\n",
-        "plt.subplot(121)\n",
-        "plt.imshow(b[:, :, 0, 2].T, vmin=-2500, vmax=2500, cmap='gray')\n",
-        "plt.title('$B_z$')\n",
-        "plt.axis('off')\n",
-        "plt.subplot(122)\n",
-        "plt.imshow(vector_norm(j).sum(2).T, origin='lower', cmap='inferno')\n",
-        "plt.title('Current Density')\n",
-        "plt.axis('off')\n",
-        "plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "RCSKCiYQSEDP"
-      },
-      "source": [
-        "By computing the potential field we can estimate the free magnetic energy, and visualize the result as integrated map. Here, we only consider the first 64 grid cells in height for faster computation."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "QvxjBEGSSXXR"
-      },
-      "outputs": [],
-      "source": [
-        "free_me = get_free_mag_energy(b[:, :, :64])\n",
-        "plt.figure(figsize=(8,4))\n",
-        "plt.imshow(free_me.sum(2).transpose(), origin='lower')\n",
-        "plt.axis('off')\n",
-        "plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ZyC2vVbza7fN"
-      },
-      "source": [
-        "## Export"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "collapsed": false,
-        "id": "Ic38DF-DTFlW"
-      },
-      "source": [
-        "VTK files can be used to visualize the extrapolation results (e.g., Paraview). We use `tvtk` for converting the files.\n",
-        "The NF2 results require little storage (about 2 MB). It is faster to download the NF2 file and convert it on your local environment using the CPU resources.\n",
-        "\n",
-        "The code below can be used to convert files in Colab."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "RnN9andvcdF8"
-      },
-      "outputs": [],
-      "source": [
-        "!pip install vtk==9.0.1\n",
-        "!pip install mayavi"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "2F0OBfVZbEMr"
-      },
-      "outputs": [],
-      "source": [
-        "from nf2.evaluation.unpack import load_cube\n",
-        "from nf2.evaluation.vtk import save_vtk"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "aphkhjLva6NK"
-      },
-      "outputs": [],
-      "source": [
-        "model_path = base_path + '/extrapolation_result.nf2'\n",
-        "vtk_path = base_path + '/extrapolation_result.vtk'\n",
-        "\n",
-        "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
-        "b = load_cube(model_path, device, progress=True)\n",
-        "save_vtk(b, vtk_path, 'B')"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "CjS3CgWqVL-d"
-      },
-      "source": [
-        "## Training - Series Extrapolation"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "K0FOZxs9YuEY"
-      },
-      "source": [
-        "Starting from our previous extrapolation, we can significantly speed-up the extrapolation of the temporal evolution. For this we adapt the model sequentially by changing the boundary condition and re-training the model. This allows us to converge within few iterations per time step."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "uWmovrrQZOsn"
-      },
-      "source": [
-        "We start by downloading the data. Fill the form below to select your time range. Note that we use the same start point and sharp as in the previous extrapolation (i.e., we require an initial extrapolation)."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Zq-PHrsSD_2i"
-      },
-      "source": [
-        "To skip the initial training you can download the pre-trained state by running the code below."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "mFsq_XVqpO4x",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Series parameters\n",
-        "step_iterations = 2e3 #@param {type:\"number\"}\n",
-        "check_val_every_n_epoch =  100 #@param {type:\"integer\"}\n",
-        "\n",
-        "series_data_args = {**data_args}\n",
-        "series_data_args['iterations'] = int(step_iterations)\n",
-        "del series_data_args['data_path']\n",
-        "\n",
-        "series_training_args = {**training_args}\n",
-        "series_training_args['lambda_b'] = 1\n",
-        "series_training_args['lr_params'] = 5e-4"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Initialize the logging to wandb."
-      ],
-      "metadata": {
-        "id": "OZdcZMV27Xkp"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "wandb.finish()\n",
-        "wandb_logger = WandbLogger(project='nf2', name=str(sharp_nr) + '_series', dir=base_path, log_model=\"all\")\n",
-        "wandb_logger.experiment.config.update(config, allow_val_change=True)"
-      ],
-      "metadata": {
-        "id": "qmlDZjwH7W31"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "We prepare the data loader which will sequentially update the boundary condtion of the magnetic field extrapolation."
-      ],
-      "metadata": {
-        "id": "AitTLhPx5vwl"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "xmEdIi04hzFF"
-      },
-      "outputs": [],
-      "source": [
-        "hmi_p_files = sorted(glob.glob(os.path.join(data_path, '*Bp.fits')))  # x\n",
-        "hmi_t_files = sorted(glob.glob(os.path.join(data_path, '*Bt.fits')))  # y\n",
-        "hmi_r_files = sorted(glob.glob(os.path.join(data_path, '*Br.fits')))  # z\n",
-        "err_p_files = sorted(glob.glob(os.path.join(data_path, '*Bp_err.fits')))  # x\n",
-        "err_t_files = sorted(glob.glob(os.path.join(data_path, '*Bt_err.fits')))  # y\n",
-        "err_r_files = sorted(glob.glob(os.path.join(data_path, '*Br_err.fits')))  # z\n",
-        "\n",
-        "data_paths = list(zip(hmi_p_files, err_p_files, hmi_t_files, err_t_files, hmi_r_files, err_r_files))\n",
-        "data_module = SHARPSeriesDataModule(data_paths, **series_data_args)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "We re-initialize the model and use the previous model as initial starting point (`meta_path`)"
-      ],
-      "metadata": {
-        "id": "VEKTlnu553iE"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "mhpsXddAh9XV"
-      },
-      "outputs": [],
-      "source": [
-        "nf2 = NF2Module(validation_settings, meta_path=save_path, **model_args, **series_training_args)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Model checkpoints are saved every epoch before we update the boundary condition. The resulting nf2 files correspond then to the frames of the series."
-      ],
-      "metadata": {
-        "id": "Zoa9McGG6GtL"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "soDvM3rNnyqF"
-      },
-      "outputs": [],
-      "source": [
-        "save_callback = LambdaCallback(on_train_epoch_end=lambda *args: save(\n",
-        "    os.path.join(base_path, os.path.basename(data_paths[nf2.current_epoch][0]).split('.')[-3] + '.nf2'),\n",
-        "    nf2.model, data_module, config, nf2.height_mapping_model))"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "xdi3DjPfhnbm"
-      },
-      "outputs": [],
-      "source": [
-        "trainer = Trainer(max_epochs=len(data_paths),\n",
-        "                  logger=wandb_logger,\n",
-        "                  devices=n_gpus if n_gpus >= 1 else None,\n",
-        "                  accelerator='gpu' if n_gpus >= 1 else None,\n",
-        "                  strategy='dp' if n_gpus > 1 else None,\n",
-        "                  num_sanity_val_steps=0, callbacks=[save_callback],\n",
-        "                  gradient_clip_val=0.1, reload_dataloaders_every_n_epochs=1,\n",
-        "                  check_val_every_n_epoch=check_val_every_n_epoch)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "r5qHAfiOp0F3"
-      },
-      "outputs": [],
-      "source": [
-        "trainer.fit(nf2, data_module)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## Evaluation  - Series"
-      ],
-      "metadata": {
-        "id": "zkV9nsGF6pO_"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "8ncM70ydD9m0",
-        "cellView": "form"
-      },
-      "outputs": [],
-      "source": [
-        "#@title Download sample extrapolation result\n",
-        "download = False #@param {type:\"boolean\"}\n",
-        "if download:\n",
-        "  zip_path = os.path.join(base_path, 'nf2_series.zip')\n",
-        "  gdown.download('https://drive.google.com/uc?id=1lkbPC7r8EiIrNp_h14Dj_DUKzwpkR60x', zip_path)\n",
-        "  shutil.unpack_archive(zip_path, base_path)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "nf2_files = sorted(glob.glob(os.path.join(base_path, '*.nf2')))\n",
-        "nf2_files = [f for f in nf2_files if 'extrapolation_result.nf2' not in f] # filter the initial state\n",
-        "\n",
-        "state = torch.load(nf2_files[0], map_location='cpu')\n",
-        "Mm_per_pix = state['Mm_per_pixel']\n",
-        "cm_per_pix = (Mm_per_pix * 1e8)\n",
-        "z_pixels = int(np.ceil(20 / (Mm_per_pix)))\n",
-        "\n",
-        "series_results = evaluate_nf2_series(nf2_files, z_pixels, cm_per_pix, 2, batch_size)"
-      ],
-      "metadata": {
-        "id": "OA-NJpAUK4Wt"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "Lby5AjxY_NO0"
-      },
-      "source": [
-        "Registered flares in this sequence: X2.2 - 08:57; C5.4 - 10:13; X9.3 - 11:53"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "KOpBr-s1_Y93"
-      },
-      "outputs": [],
-      "source": [
-        "# plot the integrated free magnetic energy\n",
-        "plt.figure(figsize=(9, 3))\n",
-        "plt.plot(series_results['date'], np.array(series_results['total_free_energy']) * 1e-32)\n",
-        "plt.ylabel('Free Energy\\n[$10^{32}$ erg]')\n",
-        "plt.show()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "## Evaluate Flare Event"
-      ],
-      "metadata": {
-        "id": "NBW3W6zcFRuL"
-      }
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "The magnetic field extrapolations can be used to compare the energy change to flare events and locate regions of strong energy release.\n",
-        "\n",
-        "We use `Fido` to find registered flares in our series."
-      ],
-      "metadata": {
-        "id": "nEMcfv6pFYoP"
-      }
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "XEMIl6UW4Ipb"
-      },
-      "outputs": [],
-      "source": [
-        "dates = [parse(os.path.basename(nf2_file).split('.')[0][:-4].replace('_', 'T')) for nf2_file in nf2_files]\n",
-        "flares = Fido.search(a.Time(min(dates), max(dates)),\n",
-        "                     a.hek.EventType(\"FL\"),\n",
-        "                     a.hek.OBS.Observatory == \"GOES\")[\"hek\"]\n",
-        "#flares = [f for f in flares if f[\"fl_goescls\"][0] in ['M', 'X']]"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "flares[['event_starttime', 'fl_goescls']]"
-      ],
-      "metadata": {
-        "id": "E88ha_WyYH0k"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Use the list above to select the flare event you are interested in."
-      ],
-      "metadata": {
-        "id": "i8IQvuABajPV"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "flare = list(flares)[3] # select flare"
-      ],
-      "metadata": {
-        "id": "AkW32-lpGW5r"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "We download EUV observations as reference for comparison."
-      ],
-      "metadata": {
-        "id": "OjyxiwLHarNU"
-      }
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "start_time = flare[\"event_starttime\"].datetime\n",
-        "end_time = flare[\"event_endtime\"].datetime\n",
-        "euv_files = download_euv(start_time, end_time, download_dir, client=client)"
-      ],
-      "metadata": {
-        "id": "1yJ-64lsaL0H"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "filter_dates = (np.array(dates) > (start_time - timedelta(minutes=12))) & \\\n",
-        "                (np.array(dates) < (end_time + timedelta(minutes=12)))\n",
-        "flare_nf2_files = np.array(nf2_files)[filter_dates]\n",
-        "\n",
-        "free_energy_0 = _calculate_free_energy(flare_nf2_files[0], z_pixels, batch_size)\n",
-        "free_energy_1 = _calculate_free_energy(flare_nf2_files[-1], z_pixels, batch_size)\n",
-        "released_energy = -np.clip(free_energy_0 - free_energy_1, a_min=None, a_max=0)\n",
-        "released_energy_map = released_energy.sum(2) * cm_per_pix\n",
-        "\n",
-        "mag_map = load_B_map(flare_nf2_files[0])\n",
-        "euv_map = get_integrated_euv_map(euv_files, mag_map.wcs)\n",
-        "\n",
-        "plt.figure(figsize=(15, 3))\n",
-        "\n",
-        "plt.subplot(131, projection=mag_map)\n",
-        "mag_map.plot()\n",
-        "plt.title('$B_z$')\n",
-        "\n",
-        "plt.subplot(132, projection=euv_map)\n",
-        "plt.imshow(euv_map.data, origin='lower', cmap='sdoaia94', norm=ImageNormalize(stretch=AsinhStretch(0.005)))\n",
-        "plt.title('Integrated EUV SDO/AIA 94 $\\AA$')\n",
-        "plt.xlabel('Carrington Longitude')\n",
-        "plt.ylabel(' ')\n",
-        "\n",
-        "plt.subplot(133, projection=mag_map)\n",
-        "im = plt.imshow(released_energy_map.T, origin='lower', cmap='jet')\n",
-        "plt.title(f'Released energy [{released_energy.sum() * cm_per_pix ** 3:.2e} erg]')\n",
-        "plt.xlabel('Carrington Longitude')\n",
-        "plt.ylabel(' ')\n",
-        "plt.colorbar(mappable=im, label='erg/cm$^2$')\n",
-        "\n",
-        "plt.tight_layout()\n",
-        "plt.show()\n",
-        "plt.close()"
-      ],
-      "metadata": {
-        "id": "sdbWfm3vF58U"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [],
-      "metadata": {
-        "id": "hBjzDuBCZSmu"
-      },
-      "execution_count": null,
-      "outputs": []
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "provenance": [],
-      "toc_visible": true,
-      "gpuType": "T4"
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python"
-    }
-  },
-  "nbformat": 4,
-  "nbformat_minor": 0
-}
\ No newline at end of file
diff --git a/nf2/convert/__init__.py b/nf2/convert/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/convert/convert_series.py b/nf2/convert/convert_series.py
new file mode 100644
index 0000000..558da89
--- /dev/null
+++ b/nf2/convert/convert_series.py
@@ -0,0 +1,58 @@
+import argparse
+import glob
+import os.path
+
+from tqdm import tqdm
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_dir', type=str, help='path to the source NF2 files')
+    parser.add_argument('--out_dir', type=str, help='path to the target VTK directory', required=False, default=None)
+    parser.add_argument('--Mm_per_pixel', type=float, help='spatial resolution (0.36 for original HMI)', required=False,
+                        default=None)
+    parser.add_argument('--height_range', type=float, nargs=2, help='height range in Mm', required=False, default=None)
+    parser.add_argument('--type', type=str, help='type of the conversion (vtk, hdf5, npy, fits)', required=False,
+                        default='vtk')
+    parser.add_argument('--overwrite', action='store_true', help='overwrite existing files', required=False,
+                        default=False)
+
+    args = parser.parse_args()
+    nf2_paths = sorted(glob.glob(args.nf2_dir))
+
+    Mm_per_pixel = args.Mm_per_pixel
+    out_dir = args.out_dir if args.out_dir is not None else os.path.dirname(args.nf2_dir)
+    os.makedirs(out_dir, exist_ok=True)
+    height_range = args.height_range
+
+    conversion_type = args.type
+    if conversion_type == 'vtk':
+        d_type = '.vtk'
+        from nf2.convert import nf2_to_vtk
+        convert_f = nf2_to_vtk.convert
+    elif conversion_type == 'hdf5':
+        d_type = '.hdf5'
+        from nf2.convert import nf2_to_hdf5
+        convert_f = nf2_to_hdf5.convert
+    elif conversion_type == 'npy':
+        d_type = '.npy'
+        from nf2.convert import nf2_to_npy
+        convert_f = nf2_to_npy.convert
+    elif conversion_type == 'fits':
+        d_type = '.fits'
+        from nf2.convert import nf2_to_fits
+        convert_f = nf2_to_fits.convert
+    else:
+        raise ValueError(f'Unknown conversion type: {conversion_type}')
+
+    for nf2_path in tqdm(nf2_paths, desc='Converting', unit='file'):
+        out_file = os.path.join(out_dir, os.path.basename(nf2_path).replace('.nf2', d_type))
+        if os.path.exists(out_file) and not args.overwrite:
+            print(f'File exists: {out_file}')
+            continue
+        convert_f(nf2_path=nf2_path, out_path=out_file,
+                  Mm_per_pixel=Mm_per_pixel, height_range=height_range, progress=False)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_fits.py b/nf2/convert/nf2_to_fits.py
new file mode 100644
index 0000000..7abc4f3
--- /dev/null
+++ b/nf2/convert/nf2_to_fits.py
@@ -0,0 +1,45 @@
+import argparse
+import os
+
+from astropy.io import fits
+
+from nf2.evaluation.output import CartesianOutput, current_density
+
+
+def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
+    out_path = out_path if out_path is not None \
+        else os.path.join(os.path.dirname(nf2_path), nf2_path.split(os.sep)[-2] + '.hdf5')
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range, **kwargs)
+
+    b = output['b']
+    j = output['j']
+
+    header = {'Mm_per_pix': output['Mm_per_pixel'],
+              'data': nf2_out.data_config,
+              'wcs': nf2_out.wcs,
+              'DATE_OBS': nf2_out.time}
+
+    b_hdu = fits.PrimaryHDU(b, header=header)
+    j_hdu = fits.FitsHDU(j)
+    hdul = fits.HDUList([b_hdu, j_hdu])
+    hdul.writeto(out_path)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to fits.')
+    parser.add_argument('nf2_path', type=str, help='path to the source NF2 file')
+    parser.add_argument('--out_path', type=str, help='path to the target numpy file', required=False, default=None)
+    parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+    strides = args.strides
+    out_path = args.out_path
+
+    convert(nf2_path, out_path, strides)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_hdf5.py b/nf2/convert/nf2_to_hdf5.py
new file mode 100644
index 0000000..ca80824
--- /dev/null
+++ b/nf2/convert/nf2_to_hdf5.py
@@ -0,0 +1,62 @@
+import argparse
+import os.path
+from threading import Thread
+
+import h5py
+
+from nf2.evaluation.output import CartesianOutput, current_density
+
+
+class _SaveFileTask(Thread):
+
+    def __init__(self, out_path, output, nf2_out):
+        super().__init__()
+        self.out_path = out_path
+        self.output = output
+        self.nf2_out = nf2_out
+    def run(self):
+        f = h5py.File(self.out_path, 'w')
+        f.create_dataset('B', data=self.output['b'], dtype='f4', compression='gzip')
+        f.attrs['INFO'] = f'Magnetic vector field (B_x, B_y, B_z): {self.output["b"].shape} (x, y, z, 3)'
+        f.attrs['Mm_per_pixel'] = self.output['Mm_per_pixel']
+        f.attrs['type'] = self.nf2_out.data_config['type']
+        f.attrs['wcs'] = self.nf2_out.wcs[0].to_header_string()
+        f.attrs['time'] = self.nf2_out.time.isoformat('T', timespec='seconds')
+        f.close()
+        print('File saved:', self.out_path)
+
+
+def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
+    out_path = out_path if out_path is not None \
+        else os.path.join(os.path.dirname(nf2_path), nf2_path.split(os.sep)[-2] + '.hdf5')
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range, **kwargs)
+
+    # save file in background
+    task = _SaveFileTask(out_path, output, nf2_out)
+    task.start()
+    return task
+
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+    parser.add_argument('--out_path', type=str, help='path to the target HDF5 file', required=False, default=None)
+    parser.add_argument('--Mm_per_pixel', type=float, help='spatial resolution (0.36 for original HMI)', required=False,
+                        default=None)
+    parser.add_argument('--height_range', type=float, nargs=2, help='height range in Mm', required=False, default=None)
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+
+    Mm_per_pixel = args.Mm_per_pixel
+    out_path = args.out_path
+    height_range = args.height_range
+
+    convert(nf2_path, out_path, Mm_per_pixel, height_range)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_npy.py b/nf2/convert/nf2_to_npy.py
new file mode 100644
index 0000000..63b0f8b
--- /dev/null
+++ b/nf2/convert/nf2_to_npy.py
@@ -0,0 +1,37 @@
+import argparse
+import os.path
+
+import numpy as np
+
+from nf2.evaluation.output import CartesianOutput
+
+
+def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
+    out_path = out_path if out_path is not None \
+        else os.path.join(os.path.dirname(nf2_path), nf2_path.split(os.sep)[-2] + '.npy')
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range, **kwargs)
+    np.save(out_path, output['b'])
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+    parser.add_argument('--out_path', type=str, help='path to the target NPY file', required=False, default=None)
+    parser.add_argument('--Mm_per_pixel', type=float, help='spatial resolution (0.36 for original HMI)', required=False,
+                        default=None)
+    parser.add_argument('--height_range', type=float, nargs=2, help='height range in Mm', required=False, default=None)
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+
+    Mm_per_pixel = args.Mm_per_pixel
+    out_path = args.out_path
+    height_range = args.height_range
+
+    convert(nf2_path, out_path, Mm_per_pixel, height_range)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_vtk.py b/nf2/convert/nf2_to_vtk.py
new file mode 100644
index 0000000..28cf17a
--- /dev/null
+++ b/nf2/convert/nf2_to_vtk.py
@@ -0,0 +1,52 @@
+import argparse
+import os.path
+from threading import Thread
+
+from nf2.evaluation.output import CartesianOutput
+from nf2.evaluation.vtk import save_vtk
+
+
+class _SaveFileTask(Thread):
+
+    def __init__(self, out_path, output):
+        super().__init__()
+        self.out_path = out_path
+        self.output = output
+
+    def run(self):
+        save_vtk(self.out_path, vectors={'B': self.output['b'], 'J': self.output['j']},
+                 Mm_per_pix=self.output['Mm_per_pixel'])
+
+
+def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
+    out_path = out_path if out_path is not None \
+        else os.path.join(os.path.dirname(nf2_path), nf2_path.split(os.sep)[-2] + '.vtk')
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range, **kwargs)
+
+    # save file in background
+    task = _SaveFileTask(out_path, output)
+    task.start()
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+    parser.add_argument('--out_path', type=str, help='path to the target VTK file', required=False, default=None)
+    parser.add_argument('--Mm_per_pixel', type=float, help='spatial resolution (0.36 for original HMI)', required=False,
+                        default=None)
+    parser.add_argument('--height_range', type=float, nargs=2, help='height range in Mm', required=False, default=None)
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+
+    Mm_per_pixel = args.Mm_per_pixel
+    out_path = args.out_path
+    height_range = args.height_range
+
+    convert(nf2_path, out_path, Mm_per_pixel, height_range)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_vtk_spherical.py b/nf2/convert/nf2_to_vtk_spherical.py
new file mode 100644
index 0000000..6d1c5e4
--- /dev/null
+++ b/nf2/convert/nf2_to_vtk_spherical.py
@@ -0,0 +1,72 @@
+import argparse
+import glob
+import os
+import threading
+
+import numpy as np
+from astropy import units as u
+
+from nf2.evaluation.metric import normalized_divergence
+from nf2.evaluation.output import SphericalOutput, current_density, twist
+from nf2.evaluation.vtk import save_vtk
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+parser.add_argument('--out_path', type=str, help='path to the target VTK files', required=False, default=None)
+parser.add_argument('--temporal_strides', type=int, default=1, required=False,
+                    help='select the time steps (default=1; all files).')
+parser.add_argument('--overwrite', action='store_true', help='overwrite existing files')
+
+parser.add_argument('--radius_range', nargs='+', type=float, default=(0.999, 1.3), required=False)
+parser.add_argument('--latitude_range', nargs='+', type=float, default=(0, 180), required=False)
+parser.add_argument('--longitude_range', nargs='+', type=float, default=(0, 360), required=False)
+parser.add_argument('--radians', action='store_true', help='latitude and longitude in radians', required=False,
+                    default=False)
+parser.add_argument('--pixels_per_solRad', type=int, default=64, required=False)
+
+args = parser.parse_args()
+nf2_path = args.nf2_path
+out_path = args.out_path if args.out_path is not None else os.path.dirname(nf2_path)
+
+os.makedirs(out_path, exist_ok=True)
+
+radius_range = tuple(args.radius_range) * u.solRad
+if args.radians:
+    latitude_range = tuple(args.latitude_range) * u.rad
+    longitude_range = tuple(args.longitude_range) * u.rad
+else:
+    latitude_range = tuple(args.latitude_range) * u.deg
+    longitude_range = tuple(args.longitude_range) * u.deg
+pixels_per_solRad = args.pixels_per_solRad * u.pix / u.solRad
+
+assert len(radius_range) == 2, 'radius_range must be a tuple of length 2'
+assert len(latitude_range) == 2, 'latitude_range must be a tuple of length 2'
+assert len(longitude_range) == 2, 'longitude_range must be a tuple of length 2'
+
+files = sorted(glob.glob(nf2_path))[::args.temporal_strides]  # if os.path.isdir(nf2_path) else [nf2_path]
+
+save_threads = []
+for i, f in enumerate(files):
+
+    vtk_path = os.path.join(out_path, os.path.basename(f).replace('.nf2', '.vtk'))
+    if os.path.exists(vtk_path) and not args.overwrite:
+        print(f'Skipping existing file {f}')
+        continue
+    print('Processing file {} of {}'.format(i + 1, len(files)))
+
+    output = SphericalOutput(f)
+    result = output.load(radius_range, latitude_range, longitude_range, pixels_per_solRad, progress=True,
+                         metrics={'j': current_density, 'twist': twist})
+
+    vectors = {'B': result['b'], 'B_rtp': result['b_rtp']}
+    radius = result['spherical_coords'][..., 0]
+    scalars = {'radius': radius, 'current_density': np.sum(result['j'] ** 2, -1) ** 0.5, 'twist': result['twist']}
+    coords = result['coords']
+
+    args = (vtk_path, coords, vectors, scalars)
+    save_thread = threading.Thread(target=save_vtk, args=args)
+    save_thread.start()
+    save_threads.append(save_thread)
+
+# wait for threats to finish
+[t.join() for t in save_threads]
diff --git a/nf2/data/combine_sst_hmi.py b/nf2/data/combine_sst_hmi.py
new file mode 100644
index 0000000..f8cc4ac
--- /dev/null
+++ b/nf2/data/combine_sst_hmi.py
@@ -0,0 +1,165 @@
+import copy
+
+import numpy as np
+from astropy.coordinates import SkyCoord
+from astropy.io import fits
+from matplotlib import pyplot as plt
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map, make_fitswcs_header
+from astropy import units as u
+
+hmi_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/7310/hmi.sharp_cea_720s.7310.20180930_092400_TAI.Br.fits')
+
+data = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+meta = fits.getheader('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+meta['crota2'] = float(meta['crota2'])
+meta['CTYPE1'] = 'CRLN-CEA'
+meta['CTYPE2'] = 'CRLT-CEA'
+meta['CRVAL1'] = meta['CRVAL1'] - .1
+meta['CRVAL2'] = meta['CRVAL2']
+
+sst_map = Map(data[2], meta)
+
+# convert to Helioprojective
+hmi_coords = all_coordinates_from_map(hmi_map).transform_to(frames.Helioprojective)
+sst_coords = all_coordinates_from_map(sst_map).transform_to(frames.Helioprojective)
+
+lims = np.array([
+    [min(hmi_coords.Tx.min().value, sst_coords.Tx.min().value), max(hmi_coords.Tx.max().value, sst_coords.Tx.max().value)],
+    [min(hmi_coords.Ty.min().value, sst_coords.Ty.min().value), max(hmi_coords.Ty.max().value, sst_coords.Ty.max().value)]])
+
+
+bottom_left = SkyCoord(lims[0, 0] * u.arcsec, lims[1, 0] * u.arcsec, frame=frames.Helioprojective, observer=sst_map.observer_coordinate)
+top_right = SkyCoord(lims[0, 1] * u.arcsec, lims[1, 1] * u.arcsec, frame=frames.Helioprojective, observer=sst_map.observer_coordinate)
+
+bottom_left_pixel = sst_map.world_to_pixel(bottom_left)
+top_right_pixel = sst_map.world_to_pixel(top_right)
+
+
+shape = (top_right_pixel.y - bottom_left_pixel.y, top_right_pixel.x - bottom_left_pixel.x)
+shape = (int(shape[0].value), int(shape[1].value))
+
+target_wcs = copy.deepcopy(sst_map.wcs)
+target_wcs.array_shape = shape
+
+ref_pixel = sst_map.reference_pixel
+ref_bottom_left = sst_map.world_to_pixel(sst_map.bottom_left_coord)
+ref_top_right = sst_map.world_to_pixel(sst_map.top_right_coord)
+target_wcs.wcs.crpix = [ref_pixel.x.value + (ref_bottom_left.x.value - bottom_left_pixel.x.value),
+                 ref_pixel.y.value + (ref_bottom_left.y.value - bottom_left_pixel.y.value)]
+
+#### SAVE FITS ####
+
+mask = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea_mask.fits')
+
+data = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_8542_2018-09-30_mosaic_Bvec_cea.fits')
+data[:, mask == 0] = np.nan
+meta = fits.getheader('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_8542_2018-09-30_mosaic_Bvec_cea.fits')
+meta['crota2'] = float(meta['crota2'])
+meta['CTYPE1'] = 'CRLN-CEA'
+meta['CTYPE2'] = 'CRLT-CEA'
+meta['CRVAL1'] = meta['CRVAL1'] - .1
+meta['CRVAL2'] = meta['CRVAL2']
+for i, l in enumerate(['p', 't', 'r']):
+    sst_map = Map(data[i], meta).reproject_to(target_wcs)
+    sst_map.save(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_B{l}_cea.fits', overwrite=True)
+    sst_map.plot(vmin=-1000, vmax=1000)
+    plt.savefig(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_8542_2018-09-30_B{l}_cea.jpg')
+    plt.close()
+
+
+data = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+data[:, mask == 0] = np.nan
+meta = fits.getheader('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+meta['crota2'] = float(meta['crota2'])
+meta['CTYPE1'] = 'CRLN-CEA'
+meta['CTYPE2'] = 'CRLT-CEA'
+meta['CRVAL1'] = meta['CRVAL1'] - .1
+meta['CRVAL2'] = meta['CRVAL2']
+for i, l in enumerate(['p', 't', 'r']):
+    sst_map = Map(data[i], meta).reproject_to(target_wcs)
+    sst_map.save(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_B{l}_cea.fits', overwrite=True)
+    sst_map.plot(vmin=-1000, vmax=1000)
+    plt.savefig(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_6173_2018-09-30_B{l}_cea.jpg')
+    plt.close()
+
+
+for i, l in enumerate(['p', 't', 'r']):
+    hmi_map = Map(f'/gpfs/gpfs0/robert.jarolim/data/nf2/7310/hmi.sharp_cea_720s.7310.20180930_092400_TAI.B{l}.fits')
+    hmi_map = hmi_map.reproject_to(target_wcs)
+    hmi_map.data[~np.isnan(sst_map.data)] = np.nan
+    hmi_map.save(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_B{l}_cea.fits', overwrite=True)
+    hmi_map.plot(vmin=-1000, vmax=1000)
+    plt.savefig(f'/gpfs/gpfs0/robert.jarolim/data/nf2/sst/corrected_hmi_2018-09-30_B{l}_cea.jpg')
+    plt.close()
+
+# PLOT overview
+hmi_p_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/7310/hmi.sharp_cea_720s.7310.20180930_092400_TAI.Bp.fits')
+hmi_t_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/7310/hmi.sharp_cea_720s.7310.20180930_092400_TAI.Bt.fits')
+hmi_r_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/7310/hmi.sharp_cea_720s.7310.20180930_092400_TAI.Br.fits')
+
+data = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+meta = fits.getheader('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea.fits')
+meta['crota2'] = float(meta['crota2'])
+meta['CTYPE1'] = 'CRLN-CEA'
+meta['CTYPE2'] = 'CRLT-CEA'
+meta['CRVAL1'] = meta['CRVAL1'] - .1
+meta['CRVAL2'] = meta['CRVAL2']
+sst_6173_p_map = Map(data[0], meta)
+sst_6173_t_map = Map(data[1], meta)
+sst_6173_r_map = Map(data[2], meta)
+
+
+data = fits.getdata('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_8542_2018-09-30_mosaic_Bvec_cea.fits')
+meta = fits.getheader('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_8542_2018-09-30_mosaic_Bvec_cea.fits')
+meta['crota2'] = float(meta['crota2'])
+meta['CTYPE1'] = 'CRLN-CEA'
+meta['CTYPE2'] = 'CRLT-CEA'
+meta['CRVAL1'] = meta['CRVAL1'] - .1
+meta['CRVAL2'] = meta['CRVAL2']
+sst_8542_p_map = Map(data[0], meta)
+sst_8542_t_map = Map(data[1], meta)
+sst_8542_r_map = Map(data[2], meta)
+
+
+fig, axs = plt.subplots(3, 3, figsize=(15, 15))
+
+axs[0, 0].imshow(hmi_p_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                 extent=[hmi_coords.Tx.min().value, hmi_coords.Tx.max().value, hmi_coords.Ty.min().value, hmi_coords.Ty.max().value])
+
+axs[0, 1].imshow(hmi_t_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[hmi_coords.Tx.min().value, hmi_coords.Tx.max().value, hmi_coords.Ty.min().value, hmi_coords.Ty.max().value])
+
+axs[0, 2].imshow(hmi_r_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[hmi_coords.Tx.min().value, hmi_coords.Tx.max().value, hmi_coords.Ty.min().value, hmi_coords.Ty.max().value])
+
+axs[1, 0].imshow(sst_6173_p_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+axs[1, 1].imshow(sst_6173_t_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+axs[1, 2].imshow(sst_6173_r_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+axs[2, 0].imshow(sst_8542_p_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+axs[2, 1].imshow(sst_8542_t_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+axs[2, 2].imshow(sst_8542_r_map.data, origin='lower', cmap='gray', vmin=-1500, vmax=1500,
+                    extent=[sst_coords.Tx.min().value, sst_coords.Tx.max().value, sst_coords.Ty.min().value, sst_coords.Ty.max().value])
+
+[ax.set_xlim(lims[0, 0], lims[0, 1]) for ax in axs.flatten()]
+[ax.set_ylim(lims[1, 0], lims[1, 1]) for ax in axs.flatten()]
+
+plt.tight_layout()
+plt.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/sst/overview.jpg')
+plt.close()
+
+# padded map
+bottom_left_coord = SkyCoord(lims[0] * u.arcsec, lims[2] * u.arcsec, frame=hmi_coords.frame)
+top_right_coord = SkyCoord(lims[1] * u.arcsec, lims[3] * u.arcsec, frame=hmi_coords.frame)
+hmi_map = hmi_map.submap(bottom_left=bottom_left_coord, top_right=top_right_coord)
+
diff --git a/nf2/data/convert_coordinates.py b/nf2/data/convert_coordinates.py
new file mode 100644
index 0000000..cf52a5d
--- /dev/null
+++ b/nf2/data/convert_coordinates.py
@@ -0,0 +1,19 @@
+import argparse
+
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+from sunpy.coordinates import frames
+from sunpy.map import Map
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--fits_file', type=str, required=True)
+parser.add_argument('--Tx', type=float, required=True)
+parser.add_argument('--Ty', type=float, required=True)
+
+args = parser.parse_args()
+
+s_map = Map(args.fits_file)
+hpc_coords = SkyCoord(args.Tx * u.arcsec, args.Ty * u.arcsec, frame=s_map.coordinate_frame)
+hcc_coords = hpc_coords.transform_to(frames.HeliographicCarrington)
+
+print(f'x: {hcc_coords.lon.value}, y: {hcc_coords.lat.value}')
diff --git a/nf2/data/dataset.py b/nf2/data/dataset.py
index fc0c11f..6dc0d02 100644
--- a/nf2/data/dataset.py
+++ b/nf2/data/dataset.py
@@ -1,20 +1,26 @@
+import os
+
 import numpy as np
 import torch
-from torch.utils.data import Dataset, IterableDataset
+from torch.utils.data import Dataset
 
+from nf2.data.util import spherical_to_cartesian
+from astropy import units as u
 
 class BatchesDataset(Dataset):
 
-    def __init__(self, batches_file_paths, batch_size):
+    def __init__(self, batches_file_paths, batch_size=2 ** 13, **kwargs):
         """Data set for lazy loading a pre-batched numpy data array.
 
         :param batches_path: path to the numpy array.
         """
         self.batches_file_paths = batches_file_paths
-        self.batch_size = batch_size
+        self.batch_size = int(batch_size)
 
     def __len__(self):
-        return np.ceil(np.load(list(self.batches_file_paths.values())[0], mmap_mode='r').shape[0] / self.batch_size).astype(np.int32)
+        ref_file = list(self.batches_file_paths.values())[0]
+        n_batches = np.ceil(np.load(ref_file, mmap_mode='r').shape[0] / self.batch_size)
+        return n_batches.astype(np.int32)
 
     def __getitem__(self, idx):
         # lazy load data
@@ -22,6 +28,9 @@ def __getitem__(self, idx):
                 for k, bf in self.batches_file_paths.items()}
         return data
 
+    def clear(self):
+        [os.remove(f) for f in self.batches_file_paths.values()]
+
 class ImageDataset(Dataset):
 
     def __init__(self, cube_shape, norm, z=0):
@@ -45,11 +54,17 @@ def __getitem__(self, idx):
 
 class CubeDataset(Dataset):
 
-    def __init__(self, cube_shape, spatial_norm, strides=1, batch_size=1024):
-        coords = np.stack(np.mgrid[:cube_shape[0]:strides, :cube_shape[1]:strides, :cube_shape[2]:strides], -1)
+    def __init__(self, coord_range, ds_per_pixel=1 / 128, batch_size=2 ** 13):
+        x_resolution = int((coord_range[0, 1] - coord_range[0, 0]) / ds_per_pixel)
+        y_resolution = int((coord_range[1, 1] - coord_range[1, 0]) / ds_per_pixel)
+        z_resolution = int((coord_range[2, 1] - coord_range[2, 0]) / ds_per_pixel)
+        coords = np.stack(np.meshgrid(np.linspace(coord_range[0, 0], coord_range[0, 1], x_resolution, dtype=np.float32),
+                                      np.linspace(coord_range[1, 0], coord_range[1, 1], y_resolution, dtype=np.float32),
+                                      np.linspace(coord_range[2, 0], coord_range[2, 1], z_resolution, dtype=np.float32),
+                                      indexing='ij'), -1)
         self.coords_shape = coords.shape[:-1]
-        coords = torch.tensor(coords / spatial_norm, dtype=torch.float32)
-        coords = coords.view((-1, 3))
+        coords = coords.reshape(-1, 3)
+        batch_size = int(batch_size)
         self.coords = np.split(coords, np.arange(batch_size, len(coords), batch_size))
 
     def __len__(self):
@@ -57,33 +72,155 @@ def __len__(self):
 
     def __getitem__(self, idx):
         coord = self.coords[idx]
-        return coord
+        coord = torch.tensor(coord, dtype=torch.float32)
+        return {'coords': coord}
+
+
+class RandomSphericalCoordinateDataset(Dataset):
+
+    def __init__(self, radius_range, batch_size, Mm_per_ds,
+                 latitude_range=(0, np.pi), longitude_range=(0, 2 * np.pi),
+                 radial_weighted_sampling=False, latitude_weighted_sampling=False, **kwargs):
+        self.radius_range = radius_range
+        self.Mm_per_ds = Mm_per_ds
+        self.latitude_range = latitude_range
+        self.longitude_range = longitude_range
+        self.batch_size = batch_size
+        self.float_tensor = torch.FloatTensor
+        self.radial_weighted_sampling = radial_weighted_sampling
+        self.latitude_weighted_sampling = latitude_weighted_sampling
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, item):
+        random_coords = self.float_tensor(self.batch_size, 3).uniform_()
+        # r [1, height]
+        h_r = self.radius_range
+        if self.radial_weighted_sampling:
+            v_min, v_max = np.min(np.log(h_r)), np.max(np.log(h_r))
+            random_coords[:, 0] = v_min + random_coords[:, 0] * (v_max - v_min)
+            random_coords[:, 0] = torch.exp(random_coords[:, 0])
+        else:
+            random_coords[:, 0] = h_r[0] + random_coords[:, 0] * (h_r[1] - h_r[0])
+        # theta [0, pi]
+        if self.latitude_weighted_sampling:
+            lat_r = self.latitude_range
+            v_min, v_max = np.min(np.cos(lat_r)), np.max(np.cos(lat_r))
+            random_coords[:, 1] = v_min + random_coords[:, 1] * (v_max - v_min)
+            random_coords[:, 1] = torch.arccos(random_coords[:, 1])
+        else:
+            lat_r = self.latitude_range
+            random_coords[:, 1] = lat_r[0] + random_coords[:, 1] * (lat_r[1] - lat_r[0])
+        # phi [0, 2pi]
+        lon_r = self.longitude_range
+        random_coords[:, 2] = lon_r[0] + random_coords[:, 2] * (lon_r[1] - lon_r[0])
+        # convert to cartesian
+        random_coords = spherical_to_cartesian(random_coords, f=torch)
+        random_coords = random_coords * (1 * u.solRad).to_value(u.Mm) / self.Mm_per_ds
+        return {'coords': random_coords}
+
+
+class SphereDataset(Dataset):
+
+    def __init__(self, radius_range, Mm_per_ds, resolution=256, batch_size=1024, latitude_range=(0, np.pi), longitude_range=(0, 2 * np.pi), **kwargs):
+        ratio = (latitude_range[1] - latitude_range[0]) / (longitude_range[1] - longitude_range[0])
+        resolution_lat = int(resolution * ratio)
+        coords = np.stack(
+            np.meshgrid(np.linspace(radius_range[0], radius_range[1], resolution),
+                        np.linspace(latitude_range[0], latitude_range[1], resolution_lat),
+                        np.linspace(longitude_range[0], longitude_range[1], resolution),
+                        indexing='ij')).T
+        self.coords_shape = coords.shape[:-1]
 
+        coords = spherical_to_cartesian(coords)
+        coords = torch.tensor(coords, dtype=torch.float32)
+        coords = coords.reshape((-1, 3))
+        coords = coords * (1 * u.solRad).to_value(u.Mm) / Mm_per_ds
+        self.coords = np.split(coords, np.arange(batch_size, len(coords), batch_size))
+
+    def __len__(self):
+        return len(self.coords)
+
+    def __getitem__(self, idx):
+        coord = self.coords[idx]
+        return {'coords': coord}
+
+class SphereSlicesDataset(Dataset):
+
+    def __init__(self, radius_range, Mm_per_ds, latitude_range=(0, np.pi), longitude_range=(0, 2 * np.pi), longitude_resolution=256, batch_size=1024, n_slices=5, **kwargs):
+        ratio = (latitude_range[1] - latitude_range[0]) / (longitude_range[1] - longitude_range[0])
+        resolution_lat = int(longitude_resolution * ratio)
+        coords = np.stack(
+            np.meshgrid(np.linspace(radius_range[0], radius_range[1], n_slices),
+                        np.linspace(latitude_range[0], latitude_range[1], resolution_lat),
+                        np.linspace(longitude_range[0], longitude_range[1], longitude_resolution),
+                        indexing='ij'), -1)
+        self.cube_shape = coords.shape[:-1]
+
+        coords = spherical_to_cartesian(coords)
+        coords = torch.tensor(coords, dtype=torch.float32)
+        coords = coords.reshape((-1, 3))
+        coords = coords * (1 * u.solRad).to_value(u.Mm) / Mm_per_ds
+        self.coords = np.split(coords, np.arange(batch_size, len(coords), batch_size))
+
+    def __len__(self):
+        return len(self.coords)
+
+    def __getitem__(self, idx):
+        coord = self.coords[idx]
+        return {'coords': coord}
+
+class SlicesDataset(Dataset):
+
+    def __init__(self, coord_range, ds_per_pixel, n_slices=10, batch_size=4096, **kwargs):
+        x_resolution = int((coord_range[0, 1] - coord_range[0, 0]) / ds_per_pixel)
+        y_resolution = int((coord_range[1, 1] - coord_range[1, 0]) / ds_per_pixel)
+        coords = np.stack(np.meshgrid(np.linspace(coord_range[0, 0], coord_range[0, 1], x_resolution, dtype=np.float32),
+                                        np.linspace(coord_range[1, 0], coord_range[1, 1], y_resolution, dtype=np.float32),
+                                        np.linspace(coord_range[2, 0], coord_range[2, 1], n_slices, dtype=np.float32),
+                                        indexing='ij'), -1)
+        self.cube_shape = coords.shape[:-1]
+        #
+        coords = coords.reshape((-1, 3))
+        self.coords = np.split(coords, np.arange(batch_size, len(coords), batch_size))
+
+        super().__init__()
+
+    def __len__(self):
+        return len(self.coords)
+
+    def __getitem__(self, idx):
+        coord = self.coords[idx]
+        coord = torch.tensor(coord, dtype=torch.float32)
+        return {'coords': coord}
 
 class RandomCoordinateDataset(Dataset):
 
-    def __init__(self, cube_shape, spatial_norm, batch_size, buffer=None):
+    def __init__(self, coord_range, batch_size=2 ** 14, buffer=None, z_sampling_exponent=1):
         super().__init__()
-        cube_shape = np.array([[0, cube_shape[0] - 1], [0, cube_shape[1] - 1], [0, cube_shape[2] - 1]])
         if buffer:
-            buffer_x = (cube_shape[0, 1] - cube_shape[0, 0]) * buffer
-            buffer_y = (cube_shape[1, 1] - cube_shape[1, 0]) * buffer
-            cube_shape[0, 0] -= buffer_x
-            cube_shape[0, 1] += buffer_x
-            cube_shape[1, 0] -= buffer_y
-            cube_shape[1, 1] += buffer_y
-        self.cube_shape = cube_shape
-        self.spatial_norm = spatial_norm
-        self.batch_size = batch_size
+            buffer_x = (coord_range[0, 1] - coord_range[0, 0]) * buffer
+            buffer_y = (coord_range[1, 1] - coord_range[1, 0]) * buffer
+            coord_range[0, 0] -= buffer_x
+            coord_range[0, 1] += buffer_x
+            coord_range[1, 0] -= buffer_y
+            coord_range[1, 1] += buffer_y
+        self.coord_range = coord_range
+        self.batch_size = int(batch_size)
         self.float_tensor = torch.FloatTensor
+        self.z_sampling_exponent = torch.tensor(z_sampling_exponent, dtype=torch.float32)
 
     def __len__(self):
         return 1
 
     def __getitem__(self, item):
         random_coords = self.float_tensor(self.batch_size, 3).uniform_()
-        random_coords[:, 0] = (random_coords[:, 0] * (self.cube_shape[0, 1] - self.cube_shape[0, 0]) + self.cube_shape[0, 0])
-        random_coords[:, 1] = (random_coords[:, 1] * (self.cube_shape[1, 1] - self.cube_shape[1, 0]) + self.cube_shape[1, 0])
-        random_coords[:, 2] = (random_coords[:, 2] * (self.cube_shape[2, 1] - self.cube_shape[2, 0]) + self.cube_shape[2, 0])
-        random_coords /= self.spatial_norm
-        return random_coords
\ No newline at end of file
+        random_coords[:, 0] = (
+                    random_coords[:, 0] * (self.coord_range[0, 1] - self.coord_range[0, 0]) + self.coord_range[0, 0])
+        random_coords[:, 1] = (
+                    random_coords[:, 1] * (self.coord_range[1, 1] - self.coord_range[1, 0]) + self.coord_range[1, 0])
+        random_coords[:, 2] = random_coords[:, 2] ** self.z_sampling_exponent
+        random_coords[:, 2] = (
+                    random_coords[:, 2] * (self.coord_range[2, 1] - self.coord_range[2, 0]) + self.coord_range[2, 0])
+        return {'coords': random_coords}
diff --git a/nf2/data/download.py b/nf2/data/download.py
index c3faa5c..54fc5ed 100644
--- a/nf2/data/download.py
+++ b/nf2/data/download.py
@@ -1,41 +1,57 @@
 import os
 
+import drms
+from dateutil.parser import parse
 
-def download_HARP(harpnum, time, dir, client, series='sharp_cea_720s', download_error=True):
-    segments = 'Br, Bp, Bt, Br_err, Bp_err, Bt_err' if download_error else 'Br, Bp, Bt'
+
+def download_HARP(harpnum, time, dir, client, series='sharp_cea_720s', segments='Br, Bp, Bt, Br_err, Bp_err, Bt_err'):
     ds = 'hmi.%s[%d][%s]{%s}' % (
-    series, harpnum, time.isoformat('_', timespec='seconds'), segments)
+        series, harpnum, time.isoformat('_', timespec='seconds'), segments)
     donwload_ds(ds, dir, client)
 
-def download_HARP_series(harpnum, t_start, duration, download_dir, client, series='sharp_cea_720s', download_error=True):
-    segments = 'Br, Bp, Bt, Br_err, Bp_err, Bt_err' if download_error else 'Br, Bp, Bt'
-    ds = 'hmi.%s[%d][%s/%s]{%s}' % \
-         (series, harpnum, t_start.isoformat('_', timespec='seconds'), duration, segments)
-    donwload_ds(ds, download_dir, client)
 
-def donwload_ds(ds, dir, client):
+def download_SHARP_series(download_dir, email, t_start, t_end=None, noaa_num=None, sharp_num=None,
+             cadence='720s', segments='Br, Bp, Bt, Br_err, Bp_err, Bt_err', series='sharp_cea_720s'):
+    assert sharp_num is not None or noaa_num is not None, 'Either sharp_num or noaa_num must be provided'
+    os.makedirs(download_dir, exist_ok=True)
+    client = drms.Client(email=email, verbose=True)
+    if noaa_num is not None:
+        sharp_num = find_HARP(t_start, noaa_num, client)
+        print(f'Found HARP number {sharp_num} for NOAA number {noaa_num}')
+    else:
+        sharp_num = sharp_num
+
+    if t_end is None:
+        ds = f'hmi.{series}[{sharp_num}][{t_start.isoformat("_", timespec="seconds")}]{{{segments}}}'
+    else:
+        duration = (t_end - t_start).total_seconds()
+        ds = f'hmi.{series}[{sharp_num}][{t_start.isoformat("_", timespec="seconds")}/{duration}s@{cadence}]{{{segments}}}'
+    return donwload_ds(ds, download_dir, client)
+
+
+def donwload_ds(ds, dir, client, process=None):
     os.makedirs(dir, exist_ok=True)
-    r = client.export(ds, protocol='fits')
+    r = client.export(ds, protocol='fits', process=process)
     r.wait()
     download_result = r.download(dir)
     return download_result
 
 
-def find_HARP(start_time, noaa_nums, client):
+def find_HARP(start_time, noaa_num, client):
     ar_mapping = client.query('hmi.Mharp_720s[][%sZ]' % start_time.isoformat('_', timespec='seconds'),
                               key=['NOAA_AR', 'HARPNUM'])
     if len(ar_mapping) == 0:
         return None
-    if noaa_nums is None or len(noaa_nums) == 0:
-        return {'HARP': list(ar_mapping['HARPNUM']), 'NOAA': list(ar_mapping['NOAA_AR'])}
-    for noaa_num in noaa_nums:
-        harpnum = ar_mapping[ar_mapping['NOAA_AR'] == int(noaa_num)]['HARPNUM']
-        if len(harpnum) > 0:
-            return harpnum.iloc[0]
+    harpnum = ar_mapping[ar_mapping['NOAA_AR'] == int(noaa_num)]['HARPNUM']
+    if len(harpnum) > 0:
+        return harpnum.iloc[0]
     return None
 
 
-def download_euv(start_time, end_time, dir, client):
-    ds = f'aia.lev1_euv_12s[{start_time.isoformat("_", timespec="seconds")} / {(end_time - start_time).total_seconds()}s@60s][94]{{image}}'
+def download_euv(start_time, dir, client, end_time=None, cadence='60s', channel='131'):
+    if end_time is None:
+        ds = f'aia.lev1_euv_12s[{start_time.isoformat("_", timespec="seconds")}][{channel}]{{image}}'
+    else:
+        ds = f'aia.lev1_euv_12s[{start_time.isoformat("_", timespec="seconds")} / {(end_time - start_time).total_seconds()}s@{cadence}][{channel}]{{image}}'
     euv_files = donwload_ds(ds, dir, client).download
-    return euv_files
\ No newline at end of file
+    return euv_files
diff --git a/nf2/data/download_euv.py b/nf2/data/download_euv.py
new file mode 100644
index 0000000..b69c3fd
--- /dev/null
+++ b/nf2/data/download_euv.py
@@ -0,0 +1,26 @@
+import argparse
+import os
+
+import drms
+from dateutil.parser import parse
+
+from nf2.data.download import download_euv
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--download_dir', type=str, required=True)
+parser.add_argument('--email', type=str, required=True)
+parser.add_argument('--t_start', type=str, required=True)
+parser.add_argument('--t_end', type=str, required=False, default=None)
+parser.add_argument('--cadence', type=str, required=False, default='60s')
+parser.add_argument('--channel', type=str, required=False, default='131')
+args = parser.parse_args()
+
+os.makedirs(args.download_dir, exist_ok=True)
+client = drms.Client(email=(args.email), verbose=True)
+
+start_time = parse(args.t_start)
+end_time = parse(args.t_end) if args.t_end is not None else None
+
+download_euv(start_time=start_time, end_time=end_time,
+             cadence=args.cadence, channel=args.channel,
+             dir=args.download_dir, client=client)
diff --git a/nf2/data/download_full_disk.py b/nf2/data/download_full_disk.py
new file mode 100644
index 0000000..df31fa2
--- /dev/null
+++ b/nf2/data/download_full_disk.py
@@ -0,0 +1,49 @@
+import argparse
+import glob
+import os
+
+import drms
+from dateutil.parser import parse
+
+from nf2.data.download import download_HARP, donwload_ds
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--download_dir', type=str, required=True)
+parser.add_argument('--email', type=str, required=True)
+parser.add_argument('--carrington_rotation', type=int, required=False, default=None)
+parser.add_argument('--t_start', type=str, required=True)
+parser.add_argument('--t_end', type=str, required=False, default=None)
+parser.add_argument('--convert_ptr', action='store_true')
+parser.add_argument('--download_synoptic', action='store_true')
+args = parser.parse_args()
+
+
+time = parse(args.t_start)
+t_end = parse(args.t_end) if args.t_end is not None else None
+
+os.makedirs(args.download_dir, exist_ok=True)
+client = drms.Client(email=args.email, verbose=True)
+
+# download synoptic data
+if args.download_synoptic:
+    assert args.carrington_rotation is not None, "Carrington rotation must be specified to download synoptic data"
+    # download corrected synoptic data
+    ds = f'hmi.synoptic_mr_polfil_720s[{args.carrington_rotation}]'
+    donwload_ds(ds, args.download_dir, client)
+
+    ds = f'hmi.b_synoptic[{args.carrington_rotation}]{{Br, Bt, Bp}}'
+    donwload_ds(ds, args.download_dir, client)
+
+# download full disk data
+segments = '' if args.convert_ptr else '{field, inclination, azimuth, disambig}'
+process = {'HmiB2ptr': {'l': 1}} if args.convert_ptr else None
+
+if t_end is None:
+    ds = f"hmi.B_720s[{time.isoformat('_', timespec='seconds')}]{segments}"
+else:
+    ds = f"hmi.B_720s[{time.isoformat('_', timespec='seconds').replace('-', '.')}-{t_end.isoformat('_', timespec='seconds').replace('-', '.')}]{segments}"
+full_disk_dir = os.path.join(args.download_dir, 'full_disk')
+donwload_ds(ds, full_disk_dir, client, process=process)
+
+[os.remove(f) for f in glob.glob(os.path.join(full_disk_dir, '*lat.fits'))]
+[os.remove(f) for f in glob.glob(os.path.join(full_disk_dir, '*lon.fits'))]
\ No newline at end of file
diff --git a/nf2/data/download_range.py b/nf2/data/download_range.py
index 2c43f42..74d205a 100644
--- a/nf2/data/download_range.py
+++ b/nf2/data/download_range.py
@@ -1,26 +1,44 @@
 import argparse
-import os
 
-import drms
 from dateutil.parser import parse
 
-from nf2.data.download import donwload_ds, download_HARP_series
+from nf2.data.download import download_SHARP_series
 
-parser = argparse.ArgumentParser()
-parser.add_argument('--download_dir', type=str, required=True)
-parser.add_argument('--email', type=str, required=True)
-parser.add_argument('--harpnum', type=int, required=True)
-parser.add_argument('--t_start', type=str, required=True)
-parser.add_argument('--duration', type=str, required=False, default='1d')
-parser.add_argument('--series', type=str, required=False, default='sharp_cea_720s')
-parser.add_argument('--no_error', action='store_false')
-args = parser.parse_args()
 
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--download_dir', type=str, required=True)
+    parser.add_argument('--email', type=str, required=True)
+    parser.add_argument('--sharp_num', type=int, required=False, default=None)
+    parser.add_argument('--noaa_num', type=int, required=False, default=None)
+    parser.add_argument('--t_start', type=str, required=True)
+    parser.add_argument('--t_end', type=str, required=False, default=None)
+    parser.add_argument('--cadence', type=str, required=False, default='720s')
+    parser.add_argument('--series', type=str, required=False, default='sharp_cea_720s')
+    parser.add_argument('--segments', type=str, required=False, default='Br, Bp, Bt, Br_err, Bp_err, Bt_err')
+    args = parser.parse_args()
 
+    sharp_num = args.harp_num
+    noaa_num = args.noaa_num
+    t_start = args.t_start
+    t_end = args.t_end
 
-os.makedirs(args.download_dir, exist_ok=True)
-client = drms.Client(email=(args.email), verbose=True)
-download_HARP_series(args.harpnum, parse(args.t_start), args.duration, args.download_dir, client, args.series, download_error=args.no_error)
+    segments = args.segments
+    series = args.series
+    cadence = args.cadence
 
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
+    download_dir = args.download_dir
+    email = args.email
+
+    t_start = parse(t_start)
+    t_end = parse(t_end) if t_end is not None else None
+
+    download_SHARP_series(download_dir, email, t_start, t_end, noaa_num, sharp_num, cadence, segments, series)
+
+
+if __name__ == '__main__':
+    main()
+
+
+def main():  # workaround for entry_points
+    pass
diff --git a/nf2/data/download_sharp.py b/nf2/data/download_sharp.py
index 4c14067..2ccb40b 100644
--- a/nf2/data/download_sharp.py
+++ b/nf2/data/download_sharp.py
@@ -12,9 +12,9 @@
 parser.add_argument('--harpnum', type=int, required=True)
 parser.add_argument('--t_start', type=str, required=True)
 parser.add_argument('--series', type=str, required=False, default='sharp_cea_720s')
-parser.add_argument('--no_error', action='store_false')
+parser.add_argument('--segments', type=str, required=False, default='Br, Bp, Bt, Br_err, Bp_err, Bt_err')
 args = parser.parse_args()
 
 os.makedirs(args.download_dir, exist_ok=True)
 client = drms.Client(email=(args.email), verbose=True)
-download_HARP(args.harpnum, parse(args.t_start), args.download_dir, client, args.series, download_error=args.no_error)
+download_HARP(args.harpnum, parse(args.t_start), args.download_dir, client, args.series, segments=args.segments)
diff --git a/nf2/data/download_synoptic.py b/nf2/data/download_synoptic.py
new file mode 100644
index 0000000..6d682cc
--- /dev/null
+++ b/nf2/data/download_synoptic.py
@@ -0,0 +1,26 @@
+import argparse
+import glob
+import os
+
+import drms
+from dateutil.parser import parse
+
+from nf2.data.download import download_HARP, donwload_ds
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--download_dir', type=str, required=True)
+parser.add_argument('--email', type=str, required=True)
+parser.add_argument('--carrington_rotation_start', type=int, required=True)
+parser.add_argument('--carrington_rotation_end', type=int, required=True)
+args = parser.parse_args()
+
+os.makedirs(args.download_dir, exist_ok=True)
+client = drms.Client(email=(args.email), verbose=True)
+
+for carrington_rotation in range(args.carrington_rotation_start, args.carrington_rotation_end + 1):
+    # download corrected synoptic data
+    ds = f'hmi.synoptic_mr_polfil_720s[{carrington_rotation}]'
+    donwload_ds(ds, args.download_dir, client)
+
+    ds = f'hmi.b_synoptic[{carrington_rotation}]{{Br, Bt, Bp}}'
+    donwload_ds(ds, args.download_dir, client)
\ No newline at end of file
diff --git a/nf2/data/get_coordinate_bounds.py b/nf2/data/get_coordinate_bounds.py
new file mode 100644
index 0000000..e63a12c
--- /dev/null
+++ b/nf2/data/get_coordinate_bounds.py
@@ -0,0 +1,31 @@
+import argparse
+
+import numpy as np
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
+import astropy.units as u
+
+# arguments
+parser = argparse.ArgumentParser(description='Determines the Carrington coordinate bounds of a FITS file.')
+
+parser.add_argument('file', type=str, help='FITS file to get the coordinate bounds of')
+parser.add_argument('--max_longitude', type=float, default=0.45 * np.pi, help='Buffer to add to the coordinate bounds')
+parser.add_argument('--max_latitude', type=float, default=0.4 * np.pi, help='Buffer to add to the coordinate bounds')
+args = parser.parse_args()
+
+s_map = Map(args.file)
+center_coord = s_map.center.transform_to(frames.HeliographicCarrington)
+center_lat = np.pi / 2 * u.rad - center_coord.lat
+center_lon = center_coord.lon
+
+lat_bounds = np.array([center_lat.to_value(u.rad) - args.max_latitude,
+                       center_lat.to_value(u.rad) + args.max_latitude])
+
+lon_bounds = np.array([center_lon.to_value(u.rad) - args.max_longitude,
+                       center_lon.to_value(u.rad) + args.max_longitude])
+if lon_bounds[0] < 0:
+    lon_bounds += 2 * np.pi
+
+print(f'Center: lat {center_lat.to(u.rad):.3f}, lon {center_lon.to(u.rad):.3f}')
+print(f'Longitude bounds: {lon_bounds[0]:.3f}, {lon_bounds[1]:.3f}')
+print(f'Latitude bounds: {lat_bounds[0]:.3f}, {lat_bounds[1]:.3f}')
diff --git a/nf2/data/loader.py b/nf2/data/loader.py
index 40ebb79..52d6068 100644
--- a/nf2/data/loader.py
+++ b/nf2/data/loader.py
@@ -7,6 +7,7 @@
 from matplotlib import pyplot as plt
 from sunpy.map import Map, all_coordinates_from_map
 
+from nf2.data.util import spherical_to_cartesian
 from nf2.potential.potential_field import get_potential_boundary, get_potential_top
 
 
@@ -36,35 +37,13 @@ def prep_b_data(b_cube, error_cube, height,
 
     return coords, values, err
 
-
-def load_spherical_hmi_data(data_path):
-    if isinstance(data_path, str):
-        hmi_p = sorted(glob.glob(os.path.join(data_path, '*Bp.fits')))[0]  # x
-        hmi_t = sorted(glob.glob(os.path.join(data_path, '*Bt.fits')))[0]  # y
-        hmi_r = sorted(glob.glob(os.path.join(data_path, '*Br.fits')))[0]  # z
-    else:
-        hmi_p, hmi_r, hmi_t = data_path
-
-    p_map = Map(hmi_p)  # use as coordinate reference
-    t_map = Map(hmi_t)
-    r_map = Map(hmi_r)
-
-    coords = all_coordinates_from_map(p_map)
-    coords = np.stack([np.deg2rad(coords.lon.value),
-                       np.pi / 2 + np.deg2rad(coords.lat.value),
-                       coords.radius.value]).transpose()
-
-    hmi_cube = np.stack([p_map.data, -t_map.data, r_map.data]).transpose()
-    return coords, hmi_cube, r_map.meta
-
-
-def load_potential_field_data(hmi_cube, height, reduce, only_top=False, pf_error=0.0, **kwargs):
+def load_potential_field_data(bz, height, reduce, only_top=False, pf_error=0.0, **kwargs):
     if reduce > 1:
-        hmi_cube = block_reduce(hmi_cube, (reduce, reduce, 1), func=np.mean)
+        bz = block_reduce(bz, (reduce, reduce), func=np.mean)
         height = height // reduce
-    pf_batch_size = int(1024 * 512 ** 2 / np.prod(hmi_cube.shape[:2]))  # adjust batch to AR size
-    pf_coords, pf_values = get_potential_top(hmi_cube[:, :, 2], height, batch_size=pf_batch_size, **kwargs) \
-        if only_top else get_potential_boundary(hmi_cube[:, :, 2], height, batch_size=pf_batch_size, **kwargs)
+    pf_batch_size = int(1024 * 512 ** 2 / np.prod(bz.shape))  # adjust batch to AR size
+    pf_coords, pf_values = get_potential_top(bz, height, batch_size=pf_batch_size, **kwargs) \
+        if only_top else get_potential_boundary(bz, height, batch_size=pf_batch_size, **kwargs)
     pf_values = np.array(pf_values, dtype=np.float32)
     pf_coords = np.array(pf_coords, dtype=np.float32) * reduce  # expand to original coordinate spacing
     pf_err = np.ones_like(pf_values) * pf_error
@@ -83,29 +62,4 @@ def _plot_data(error_cube, n_hmi_cube, plot_path, b_norm):
     axs[1].imshow(error_cube[..., 1].transpose(), vmin=0, cmap='gray', origin='lower')
     axs[2].imshow(error_cube[..., 2].transpose(), vmin=0, cmap='gray', origin='lower')
     plt.savefig(os.path.join(plot_path, 'b_err.jpg'))
-    plt.close()
-
-
-class RandomSphericalCoordinateSampler():
-
-    def __init__(self, height, batch_size, cuda=True):
-        self.height = height
-        self.batch_size = batch_size
-        self.float_tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
-
-    def load_sample(self):
-        random_coords = self.float_tensor(self.batch_size, 3).uniform_()
-        random_coords[:, 0] = random_coords[:, 0] * 2 * np.pi  # phi [0, 2pi]
-        random_coords[:, 1] = random_coords[:, 1] * np.pi  # theta [0, pi]
-        random_coords[:, 2] = 1 + random_coords[:, 2] * (self.height - 1)  # r [1, height]
-        random_coords = self.to_cartesian(random_coords)
-        return random_coords
-
-    def to_cartesian(self, v):
-        sin = torch.sin
-        cos = torch.cos
-        p, t, r = v[..., 0], v[..., 1], v[..., 2]
-        x = r * sin(t) * cos(p)
-        y = r * sin(t) * sin(p)
-        z = r * cos(t)
-        return torch.stack([x, y, z], -1)
+    plt.close()
\ No newline at end of file
diff --git a/nf2/data/sst/__init__.py b/nf2/data/sst/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/data/sst/prep.py b/nf2/data/sst/prep.py
new file mode 100644
index 0000000..ec9aacd
--- /dev/null
+++ b/nf2/data/sst/prep.py
@@ -0,0 +1,160 @@
+import glob
+import os.path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from astropy import units as u
+from sunpy.coordinates import frames
+from sunpy.map import Map, make_fitswcs_header
+
+hmi_field = Map('/glade/work/rjarolim/data/SST/sharp/hmi.b_720s.20180930_092400_TAI.field.fits')
+hmi_azi = Map('/glade/work/rjarolim/data/SST/sharp/hmi.b_720s.20180930_092400_TAI.azimuth.fits')
+hmi_inc = Map('/glade/work/rjarolim/data/SST/sharp/hmi.b_720s.20180930_092400_TAI.inclination.fits')
+hmi_disamb = Map('/glade/work/rjarolim/data/SST/sharp/hmi.b_720s.20180930_092400_TAI.disambig.fits')
+
+sst_los_f = sorted(glob.glob('/glade/work/rjarolim/data/SST/BLOS*.fits'))
+sst_trv_f = sorted(glob.glob('/glade/work/rjarolim/data/SST/BTRV*.fits'))
+sst_azi_f = sorted(glob.glob('/glade/work/rjarolim/data/SST/BAZI*.fits'))
+
+
+for f_los, f_trv, f_azi in zip(sst_los_f, sst_trv_f, sst_azi_f):
+    sst_los_map = Map(f_los)
+    sst_trv_map = Map(f_trv)
+    sst_azi_map = Map(f_azi)
+
+    ref_pix = (sst_los_map.reference_pixel.x.value, sst_los_map.reference_pixel.y.value) * u.pix
+    scale = (0.0075, 0.0075) * u.deg / u.pix
+    cea_header = make_fitswcs_header(sst_los_map.data,
+                                     sst_los_map.reference_coordinate.transform_to(frames.HeliographicCarrington),
+                                     ref_pix, scale=scale, projection_code='CEA')
+
+    cea_sst_los_map = sst_los_map.reproject_to(cea_header)
+    cea_sst_trv_map = sst_trv_map.reproject_to(cea_header)
+    cea_sst_azi_map = sst_azi_map.reproject_to(cea_header)
+
+    mask = np.isnan(cea_sst_los_map.data) | (cea_sst_azi_map.data == 0)
+
+    cea_sst_los_map.data[mask] = np.nan
+    cea_sst_trv_map.data[mask] = np.nan
+    cea_sst_azi_map.data[mask] = np.nan
+
+    cea_sst_los_map.save(os.path.join('/glade/work/rjarolim/data/SST/converted', os.path.basename(f_los)), overwrite=True)
+    cea_sst_trv_map.save(os.path.join('/glade/work/rjarolim/data/SST/converted', os.path.basename(f_trv)), overwrite=True)
+    cea_sst_azi_map.save(os.path.join('/glade/work/rjarolim/data/SST/converted', os.path.basename(f_azi)), overwrite=True)
+
+    fig = plt.figure(figsize=(10, 10))
+
+    ax = plt.subplot(3, 2, 1, projection=cea_sst_los_map)
+    sst_los_map.plot(axes=ax, cmap='gray')
+    ax.set_title('LOS')
+    plt.colorbar()
+
+    ax = plt.subplot(3, 2, 2, projection=cea_sst_los_map)
+    cea_sst_los_map.plot(axes=ax, cmap='gray')
+    ax.set_title('LOS CEA')
+    plt.colorbar()
+
+    ax = plt.subplot(3, 2, 3, projection=cea_sst_trv_map)
+    sst_trv_map.plot(axes=ax, cmap='gray')
+    ax.set_title('TRV')
+    plt.colorbar()
+
+    ax = plt.subplot(3, 2, 4, projection=cea_sst_trv_map)
+    cea_sst_trv_map.plot(axes=ax, cmap='gray')
+    ax.set_title('TRV CEA')
+    plt.colorbar()
+
+    ax = plt.subplot(3, 2, 5, projection=cea_sst_azi_map)
+    sst_azi_map.plot(axes=ax, cmap='twilight')
+    ax.set_title('AZI')
+    plt.colorbar()
+
+    ax = plt.subplot(3, 2, 6, projection=cea_sst_azi_map)
+    cea_sst_azi_map.plot(axes=ax, cmap='twilight')
+    ax.set_title('AZI CEA')
+    plt.colorbar()
+
+    fig.tight_layout()
+    fig.savefig(os.path.join('/glade/work/rjarolim/data/SST/converted/', os.path.basename(f_los).replace('.fits', '.jpg')))
+    plt.close(fig)
+
+cea_hmi_field = hmi_field.reproject_to(sst_los_map.wcs).reproject_to(cea_header)
+cea_hmi_azi = hmi_azi.reproject_to(sst_los_map.wcs).reproject_to(cea_header)
+cea_hmi_inc = hmi_inc.reproject_to(sst_los_map.wcs).reproject_to(cea_header)
+cea_hmi_disamb = hmi_disamb.reproject_to(sst_los_map.wcs).reproject_to(cea_header)
+
+field = np.nan_to_num(cea_hmi_field.data, nan=0)
+azi = np.nan_to_num(np.deg2rad(cea_hmi_azi.data), nan=0)
+inc = np.nan_to_num(np.deg2rad(cea_hmi_inc.data), nan=0)
+amb = cea_hmi_disamb.data
+
+
+mask = ~np.isnan(cea_sst_los_map.data)
+
+cea_hmi_los = Map(field * np.cos(inc), cea_header)
+cea_hmi_los_masked = Map(field * np.cos(inc), cea_header)
+cea_hmi_los_masked.data[mask] = np.nan
+cea_hmi_los_masked.save('/glade/work/rjarolim/data/SST/converted/hmi_los_masked.fits', overwrite=True)
+cea_hmi_los.save('/glade/work/rjarolim/data/SST/converted/hmi_los.fits', overwrite=True)
+
+cea_hmi_trv = Map(field * np.sin(inc), cea_header)
+cea_hmi_trv_masked = Map(field * np.sin(inc), cea_header)
+cea_hmi_trv_masked.data[mask] = np.nan
+cea_hmi_trv_masked.save('/glade/work/rjarolim/data/SST/converted/hmi_trv_masked.fits', overwrite=True)
+cea_hmi_trv.save('/glade/work/rjarolim/data/SST/converted/hmi_trv.fits', overwrite=True)
+
+amb_weak = 2
+condition = (amb.astype((int)) >> amb_weak).astype(bool)
+azi[condition] += np.pi
+
+cea_hmi_azi = Map(azi, cea_header)
+cea_hmi_azi_masked = Map(np.copy(azi), cea_header)
+cea_hmi_azi_masked.data[mask] = np.nan
+cea_hmi_azi_masked.save('/glade/work/rjarolim/data/SST/converted/hmi_azi_masked.fits', overwrite=True)
+cea_hmi_azi.save('/glade/work/rjarolim/data/SST/converted/hmi_azi.fits', overwrite=True)
+
+fig = plt.figure(figsize=(10, 10))
+
+ax = plt.subplot(4, 2, 1, projection=cea_hmi_los)
+cea_hmi_los.plot(axes=ax, cmap='gray')
+ax.set_title('LOS')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 2, projection=cea_hmi_los_masked)
+cea_hmi_los_masked.plot(axes=ax, cmap='gray')
+ax.set_title('LOS MASKED')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 3, projection=cea_hmi_trv)
+cea_hmi_trv.plot(axes=ax, cmap='gray')
+ax.set_title('TRV')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 4, projection=cea_hmi_trv_masked)
+cea_hmi_trv_masked.plot(axes=ax, cmap='gray')
+ax.set_title('TRV MASKED')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 5, projection=cea_hmi_azi)
+cea_hmi_azi.plot(axes=ax, cmap='twilight', vmin=0, vmax=2 * np.pi)
+ax.set_title('AZI')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 6, projection=cea_hmi_azi_masked)
+cea_hmi_azi_masked.plot(axes=ax, cmap='twilight', vmin=0, vmax=2 * np.pi)
+ax.set_title('AZI MASKED')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 7, projection=cea_hmi_azi)
+ax.imshow(cea_hmi_azi.data % np.pi, axes=ax, cmap='twilight', vmin=0, vmax=np.pi)
+ax.set_title('AZI AMB')
+plt.colorbar()
+
+ax = plt.subplot(4, 2, 8, projection=cea_hmi_azi_masked)
+ax.imshow(cea_hmi_azi_masked.data % np.pi, axes=ax, cmap='twilight', vmin=0, vmax=np.pi)
+ax.set_title('AZI AMB MASKED')
+plt.colorbar()
+
+fig.tight_layout()
+fig.savefig('/glade/work/rjarolim/data/SST/converted/hmi.jpg')
+plt.close(fig)
diff --git a/nf2/data/util.py b/nf2/data/util.py
new file mode 100644
index 0000000..5afd66b
--- /dev/null
+++ b/nf2/data/util.py
@@ -0,0 +1,102 @@
+import numpy as np
+
+
+def spherical_to_cartesian_matrix(c):
+    r, t, p = c[..., 0], c[..., 1], c[..., 2]
+    sin = np.sin
+    cos = np.cos
+    #
+    matrix = np.stack([
+        np.stack([sin(t) * cos(p), cos(t) * cos(p), -sin(p)], -1),
+        np.stack([sin(t) * sin(p), cos(t) * sin(p), cos(p)], -1),
+        np.stack([cos(t), -sin(t), np.zeros_like(t)], -1)
+    ], -2)
+    #
+    return matrix
+
+
+def cartesian_to_spherical_matrix(c):
+    r, t, p = c[..., 0], c[..., 1], c[..., 2]
+    sin = np.sin
+    cos = np.cos
+    #
+    matrix = np.stack([
+        np.stack([sin(t) * cos(p), sin(t) * sin(p), cos(t)], -1),
+        np.stack([cos(t) * cos(p), cos(t) * sin(p), -sin(t)], -1),
+        np.stack([-sin(p), cos(p), np.zeros_like(p)], -1)
+    ], -2)
+    #
+    return matrix
+
+
+def cartesian_rotation_matrix(theta, phi):
+    m = np.array([[np.cos(theta) * np.cos(phi), -np.sin(phi), -np.sin(theta) * np.cos(phi)],
+                  [np.cos(theta) * np.sin(phi), np.cos(phi), -np.sin(theta) * np.sin(phi)],
+                  [np.sin(theta), 0, np.cos(theta)]])
+    return m
+
+
+def vector_spherical_to_cartesian(v, c, f=np):
+    vr, vt, vp = v[..., 0], v[..., 1], v[..., 2]
+    r, t, p = c[..., 0], c[..., 1], c[..., 2]
+    sin = f.sin
+    cos = f.cos
+    #
+    vx = vr * sin(t) * cos(p) + vt * cos(t) * cos(p) - vp * sin(p)
+    vy = vr * sin(t) * sin(p) + vt * cos(t) * sin(p) + vp * cos(p)
+    vz = vr * cos(t) - vt * sin(t)
+    #
+    return f.stack([vx, vy, vz], -1)
+
+
+def vector_cartesian_to_spherical(v, c, f=np):
+    vx, vy, vz = v[..., 0], v[..., 1], v[..., 2]
+    r, t, p = c[..., 0], c[..., 1], c[..., 2]
+    sin = f.sin
+    cos = f.cos
+    #
+    vr = vx * sin(t) * cos(p) + vy * sin(t) * sin(p) + vz * cos(t)
+    vt = vx * cos(t) * cos(p) + vy * cos(t) * sin(p) - vz * sin(t)
+    vp = - vx * sin(p) + vy * cos(p)
+    #
+    return f.stack([vr, vt, vp], -1)
+
+
+def spherical_to_cartesian(v, f=np):
+    sin = f.sin
+    cos = f.cos
+    r, t, p = v[..., 0], v[..., 1], v[..., 2]
+    x = r * sin(t) * cos(p)
+    y = r * sin(t) * sin(p)
+    z = r * cos(t)
+    return f.stack([x, y, z], -1)
+
+
+def cartesian_to_spherical(v, f=np):
+    x, y, z = v[..., 0], v[..., 1], v[..., 2]
+    xy = x ** 2 + y ** 2
+
+    r = f.sqrt(xy + z ** 2)
+    nudge = (f.abs(z) < 1e-6) * 1e-6  # assure numerical stability
+    t = f.arctan2(f.sqrt(xy), z + nudge)
+    nudge = (f.abs(x) < 1e-6) * 1e-6  # assure numerical stability
+    p = f.arctan2(y, x + nudge)
+
+    return f.stack([r, t, p], -1)
+
+
+def img_to_los_trv_azi(b, f=np):
+    B_los = b[..., 2]
+    B_trv = (b[..., 0] ** 2 + b[..., 1] ** 2) ** 0.5
+    azi = f.arctan2(b[..., 0], b[..., 1])
+    b = f.stack([B_los, B_trv, azi], -1)
+    return b
+
+
+def los_trv_azi_to_img(b, ambiguous=False, f=np):
+    B_los, B_trv, azi = b[..., 0], b[..., 1], b[..., 2]
+    B_x = B_trv * f.abs(f.sin(azi)) if ambiguous else B_trv * f.sin(azi)
+    B_y = B_trv * f.abs(f.cos(azi)) if ambiguous else B_trv * f.cos(azi)
+    B_z = B_los
+    b = f.stack([B_x, B_y, B_z], -1)
+    return b
diff --git a/nf2/evaluation/compute_metrics.py b/nf2/evaluation/compute_metrics.py
new file mode 100644
index 0000000..039874f
--- /dev/null
+++ b/nf2/evaluation/compute_metrics.py
@@ -0,0 +1,30 @@
+import argparse
+
+from nf2.evaluation.metric import normalized_divergence, weighted_theta, theta_J, sigma_J
+from nf2.evaluation.output import current_density, CartesianOutput
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(progress=True, metrics={'j': current_density})
+
+    metrics = {}
+    b = output['b']
+    j = output['j']
+    metrics['normalized_divergence'] = normalized_divergence(b)
+    metrics['theta_J'] = theta_J(b, j)
+    metrics['sigma_J'] = sigma_J(b, j)
+
+    print(f'Normalized divergence: {metrics["normalized_divergence"]:.3f}')
+    print(f'Theta_J: {metrics["theta_J"]:.3f}')
+    print(f'Sigma_J (1e2): {metrics["sigma_J"] * 1e2:.3f}')
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/nf2/evaluation/convert_fits_series.py b/nf2/evaluation/convert_fits_series.py
deleted file mode 100644
index 5576db3..0000000
--- a/nf2/evaluation/convert_fits_series.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import argparse
-import glob
-import gzip
-import os
-import shutil
-
-import torch
-from tqdm import tqdm
-
-from nf2.evaluation.unpack import load_cube, save_fits
-
-parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
-parser.add_argument('nf2_path', type=str, help='path to the directory of the NF2 files')
-parser.add_argument('fits_path', type=str, help='path to the directory of output FITS files')
-parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-parser.add_argument('--height', type=int, help='maximum height of the sampled volume in pixels (default none).',
-                    required=False, default=None)
-parser.add_argument('--gz_files', action='store_true')
-
-args = parser.parse_args()
-nf2_path = args.nf2_path
-fits_path = args.fits_path
-strides = args.strides
-height = args.height
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-
-os.makedirs(fits_path, exist_ok=True)
-
-
-def _zip_paths(map_paths):
-    for map_path in map_paths:
-        with open(map_path, 'rb') as f_in, gzip.open(map_path + '.gz', 'wb') as f_out:
-            f_out.writelines(f_in)
-        os.remove(map_path)
-
-
-for f in tqdm(sorted(glob.glob(os.path.join(nf2_path, '**', '*.nf2')))):
-    b = load_cube(f, device, progress=False, strides=strides, z=height)
-    for z in range(height):
-        date_str = f.split('/')[-2][:15].replace('_', 'T')
-        meta = {'DATE-OBS': date_str, 'HEIGHT': z * 360e-3}
-        map_paths = save_fits(b[:, :, z], fits_path, '%s_%02d' % (date_str, z), meta)
-        if args.gz_files:
-            _zip_paths(map_paths)
-
-shutil.make_archive(fits_path, 'zip', fits_path)
diff --git a/nf2/evaluation/convert_vtk_series.py b/nf2/evaluation/convert_vtk_series.py
deleted file mode 100644
index 3d6c872..0000000
--- a/nf2/evaluation/convert_vtk_series.py
+++ /dev/null
@@ -1,22 +0,0 @@
-import argparse
-import glob
-import os
-
-import torch
-
-from nf2.evaluation.unpack import load_cube
-from nf2.evaluation.vtk import save_vtk
-
-parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
-parser.add_argument('nf2_path', type=str, help='path to the directory of the NF2 files')
-parser.add_argument('vtk_path', type=str, help='path to the directory of output VTK files')
-parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-
-args = parser.parse_args()
-nf2_path = args.nf2_path
-vtk_path = args.vtk_path
-strides = args.strides
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-for i, f in enumerate(sorted(glob.glob(os.path.join(nf2_path, '*.nf2')))):
-    b = load_cube(f, device, progress=True, strides=strides)
-    save_vtk(b, os.path.join(vtk_path, 'series_%04d.vtk' % i), 'B')
diff --git a/nf2/evaluation/energy.py b/nf2/evaluation/energy.py
index 50bb67a..4b60e5d 100644
--- a/nf2/evaluation/energy.py
+++ b/nf2/evaluation/energy.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from nf2.potential.potential_field import get_potential, get_potential_field
+from nf2.potential.potential_field import get_potential_field
 from nf2.evaluation.metric import energy
 
 
diff --git a/nf2/evaluation/flares.py b/nf2/evaluation/flares.py
deleted file mode 100644
index 0f899d2..0000000
--- a/nf2/evaluation/flares.py
+++ /dev/null
@@ -1,122 +0,0 @@
-import argparse
-import glob
-import os
-from datetime import timedelta
-from multiprocessing import Pool
-
-import drms
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-from astropy.visualization import ImageNormalize, AsinhStretch
-from dateutil.parser import parse
-from sunpy.map import Map
-from sunpy.net import Fido
-from sunpy.net import attrs as a
-from tqdm import tqdm
-
-from nf2.data.download import download_euv
-from nf2.evaluation.energy import get_free_mag_energy
-from nf2.evaluation.unpack import load_cube, load_B_map
-
-
-def _calculate_free_energy(nf2_file, z, batch_size):
-    b = load_cube(nf2_file, progress=False, z=z, batch_size=batch_size)
-    free_me = get_free_mag_energy(b, progress=False)
-    return free_me
-
-
-class _F:
-    def __init__(self, ref_wcs):
-        self.ref_wcs = ref_wcs
-
-    def func(self, file):
-        return Map(file).reproject_to(self.ref_wcs)
-
-
-def get_integrated_euv_map(euv_files, ref_wcs):
-    with Pool(os.cpu_count()) as p:
-        reprojected_maps = p.map(_F(ref_wcs).func, euv_files)
-    integrated_euv = np.array([m.data for m in reprojected_maps]).sum(0)
-    euv_map = Map(integrated_euv, mag_map.meta)
-    return euv_map
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
-    parser.add_argument('nf2_path', type=str, help='path to the directory of the NF2 files')
-    parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
-    parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-    parser.add_argument('--email', type=str, required=True)
-    parser.add_argument('--flare_classes', nargs='+', type=str, required=False, default=['X', 'M'])
-    args = parser.parse_args()
-
-    nf2_files = sorted(glob.glob(os.path.join(args.nf2_path, '*.nf2')))
-    result_path = args.result_path
-    os.makedirs(result_path, exist_ok=True)
-
-    batch_size = int(1e5 * torch.cuda.device_count())
-
-    # load simulation scaling
-    Mm_per_pix = .72  # torch.load(nf2_files[0])['Mm_per_pix']
-    z_pixels = int(np.ceil(20 / (Mm_per_pix)))  # 20 Mm --> pixels
-
-    # adjust scale to strides
-    Mm_per_pix *= args.strides
-    cm_per_pix = (Mm_per_pix * 1e8)
-
-    dates = [parse(os.path.basename(nf2_file).split('.')[0][:-4].replace('_', 'T')) for nf2_file in nf2_files]
-
-    flares = Fido.search(a.Time(min(dates), max(dates)),
-                         a.hek.EventType("FL"),
-                         a.hek.OBS.Observatory == "GOES")["hek"]
-    flares = [f for f in flares if f["fl_goescls"][0] in args.flare_classes]
-
-    client = drms.Client(email=args.email, verbose=True)
-
-    for flare in tqdm(flares, desc='Analyse flares'):
-        save_path = os.path.join(result_path,
-                                 f'flare_{flare["event_peaktime"].datetime.isoformat("_", timespec="seconds")}_{flare["fl_goescls"]}.jpg')
-        if os.path.exists(save_path):
-            continue
-
-        start_time = flare["event_starttime"].datetime
-        end_time = flare["event_endtime"].datetime
-        euv_files = download_euv(start_time, end_time, result_path, client=client)
-
-        filter_dates = (np.array(dates) > (start_time - timedelta(minutes=12))) & \
-                       (np.array(dates) < (end_time + timedelta(minutes=12)))
-        flare_nf2_files = np.array(nf2_files)[filter_dates]
-
-        free_energy_0 = _calculate_free_energy(flare_nf2_files[0], z_pixels, batch_size)
-        free_energy_1 = _calculate_free_energy(flare_nf2_files[-1], z_pixels, batch_size)
-        released_energy = -np.clip(free_energy_0 - free_energy_1, a_min=None, a_max=0)
-        released_energy_map = released_energy.sum(2) * cm_per_pix
-
-        mag_map = load_B_map(flare_nf2_files[0])
-        euv_map = get_integrated_euv_map(euv_files, mag_map.wcs)
-
-        plt.figure(figsize=(15, 3))
-
-        plt.subplot(131, projection=mag_map)
-        mag_map.plot()
-        plt.title('$B_z$')
-
-        plt.subplot(132, projection=euv_map)
-        plt.imshow(euv_map.data, origin='lower', cmap='sdoaia94', norm=ImageNormalize(stretch=AsinhStretch(0.005)))
-        plt.title('Integrated EUV SDO/AIA 94 $\AA$')
-        plt.xlabel('Carrington Longitude')
-        plt.ylabel(' ')
-
-        plt.subplot(133, projection=mag_map)
-        im = plt.imshow(released_energy_map.T, origin='lower', cmap='jet')
-        plt.title(f'Released energy [{released_energy.sum() * cm_per_pix ** 3:.2e} erg]')
-        plt.xlabel('Carrington Longitude')
-        plt.ylabel(' ')
-        plt.colorbar(mappable=im, label='erg/cm$^2$')
-
-        plt.tight_layout()
-        plt.savefig(save_path, dpi=300)
-        plt.close()
-
-        [os.remove(f) for f in euv_files]
diff --git a/nf2/evaluation/metric.py b/nf2/evaluation/metric.py
index c04ee18..a31d3e9 100644
--- a/nf2/evaluation/metric.py
+++ b/nf2/evaluation/metric.py
@@ -7,6 +7,12 @@ def divergence(b_field):  # (x, y, z, (xyz-field))
     div_B = np.stack([np.gradient(b_field[..., i], axis=i, edge_order=2) for i in range(3)], axis=-1).sum(-1)
     return div_B
 
+def divergence_jacobian(jac_matrix):
+    dBx_dx = jac_matrix[..., 0, 0]
+    dBy_dy = jac_matrix[..., 1, 1]
+    dBz_dz = jac_matrix[..., 2, 2]
+    div_B = dBx_dx + dBy_dy + dBz_dz
+    return div_B
 
 def curl(b_field):  # (x, y, z)
     _, dFx_dy, dFx_dz = np.gradient(b_field[..., 0], axis=[0, 1, 2], edge_order=2)
@@ -45,9 +51,24 @@ def weighted_theta(b, j=None):
     j = j if j is not None else curl(b)
     sigma = vector_norm(lorentz_force(b, j)) / vector_norm(b) / vector_norm(j)
     w_sigma = np.average((sigma), weights=vector_norm(j))
-    theta_j = np.arcsin(w_sigma) * (180 / np.pi)
+    theta_j = np.arcsin(w_sigma)
+    theta_j = np.rad2deg(theta_j)
     return theta_j
 
+def theta_J(b, j=None):
+    j = j if j is not None else curl(b)
+    norm = vector_norm(b) * vector_norm(j) + 1e-7
+    sigma = vector_norm(lorentz_force(b, j)) / norm
+    j_weight = vector_norm(j)
+    angle = np.nansum(sigma * j_weight) / np.nansum(j_weight)
+    angle = np.clip(angle, -1. + 1e-7, 1. - 1e-7)
+    t_angle = np.arcsin(angle)
+    t_angle = np.rad2deg(t_angle)
+    return t_angle
+
+def sigma_J(b, j):
+    return (vector_norm(np.cross(j, b, -1)) / vector_norm(b)).sum() / (vector_norm(j).sum() + 1e-6)
+
 
 def energy(b):
     return (b ** 2).sum(-1) / (8 * np.pi)
@@ -71,9 +92,9 @@ def evaluate(b, B):
     # result['eps_p_ll'] = (vector_norm(B[:, :, :64]) ** 2).sum() / (vector_norm(B_potential) ** 2).sum()
 
     j = curl(b)
-    result['sig_J'] = (vector_norm(np.cross(j, b, -1)) / vector_norm(b)).sum() / (vector_norm(j).sum() + 1e-6) * 1e2
+    result['sig_J'] = sigma_J(b, j) * 1e2
     J = curl(B)
-    result['sig_J_ll'] = (vector_norm(np.cross(J, B, -1)) / vector_norm(B)).sum() / (vector_norm(J).sum() + 1e-6) * 1e2
+    result['sig_J_ll'] = sigma_J(B, J) * 1e2
 
     result['L1'] = (vector_norm(np.cross(j, b, -1)) ** 2 / vector_norm(b) ** 2).mean()
     result['L2'] = (divergence(b) ** 2).mean()
@@ -84,4 +105,40 @@ def evaluate(b, B):
     result['L2n'] = (np.abs(divergence(b)) / (vector_norm(b) + 1e-8)).mean() * 1e2
     result['L2n_B'] = (np.abs(divergence(B)) / (vector_norm(B) + 1e-8)).mean() * 1e2
 
-    return result
\ No newline at end of file
+    return result
+
+
+
+
+def b_diff_error(b, B, B_error):
+    b_err_diff = np.abs(b - B)
+    b_err_diff = np.clip(b_err_diff, a_min=0, a_max=None) - B_error
+    b_err_diff = np.linalg.norm(b_err_diff, axis=-1)
+    return b_err_diff
+
+
+def b_nabla_bz(b):
+    bx = b[..., 0]
+    by = b[..., 1]
+    bz = b[..., 2]
+
+    jac_matrix = np.stack(np.gradient(b, axis=(0, 1, 2)), -1)
+    dBx_dx = jac_matrix[..., 0, 0]
+    dBx_dy = jac_matrix[..., 0, 1]
+    dBx_dz = jac_matrix[..., 0, 2]
+    dBy_dx = jac_matrix[..., 1, 0]
+    dBy_dy = jac_matrix[..., 1, 1]
+    dBy_dz = jac_matrix[..., 1, 2]
+    dBz_dx = jac_matrix[..., 2, 0]
+    dBz_dy = jac_matrix[..., 2, 1]
+    dBz_dz = jac_matrix[..., 2, 2]
+
+    norm_B = np.linalg.norm(b, axis=-1)
+
+    dnormB_dx = - norm_B ** -3 * (bx * dBx_dx + by * dBy_dx + bz * dBz_dx)
+    dnormB_dy = - norm_B ** -3 * (bx * dBx_dy + by * dBy_dy + bz * dBz_dy)
+    dnormB_dz = - norm_B ** -3 * (bx * dBx_dz + by * dBy_dz + bz * dBz_dz)
+
+    b_nabla_bz = (bx * dBz_dx + by * dBz_dy + bz * dBz_dz) / norm_B ** 2 + \
+                 (bz / norm_B) * (bx * dnormB_dx + by * dnormB_dy + bz * dnormB_dz)
+    return b_nabla_bz
\ No newline at end of file
diff --git a/nf2/evaluation/mfr/__init__.py b/nf2/evaluation/mfr/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/evaluation/muram/compare_boundary.py b/nf2/evaluation/muram/compare_boundary.py
new file mode 100644
index 0000000..078e2cd
--- /dev/null
+++ b/nf2/evaluation/muram/compare_boundary.py
@@ -0,0 +1,64 @@
+import glob
+import os
+
+import numpy as np
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.evaluation.output import HeightTransformOutput
+from nf2.loader.muram import MURaMSnapshot
+import argparse
+
+from astropy import units as u
+
+parser = argparse.ArgumentParser(description='Evaluate MURaM snapshot.')
+parser.add_argument('--nf2_path', type=str, help='path to the nf2 project.')
+parser.add_argument('--source_path', type=str, help='path to the MURaM simulation.')
+parser.add_argument('--iteration', type=int, help='iteration of the snapshot.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+nf2_out = HeightTransformOutput(args.nf2_path)
+
+nf2_heights = list(nf2_out.load_height_mapping())
+
+snapshot = MURaMSnapshot(args.source_path, args.iteration)
+tau_cube = snapshot.tau
+
+pix_height = np.argmin(np.abs(tau_cube - 1), axis=2) * u.pix
+Mm_height = (pix_height * snapshot.ds[2]).to(u.Mm)
+base_height_Mm = Mm_height.min()
+base_height_pix = pix_height.min()
+cube_height = (snapshot.shape[2] * u.pix * snapshot.ds[2]).to(u.Mm)
+print(f'Cube Height: {cube_height} [Mm]; {snapshot.shape[2]} [pix]')
+print(f'Base Height: {base_height_Mm} [Mm]; {base_height_pix} [pix]')
+
+fig, axs = plt.subplots(2, 4, figsize=(16, 4))
+
+for i, (tau, nf2_height) in enumerate(zip([1.0, 0.001, 0.000100, 0.000001], nf2_heights)):
+    row = axs[:, i]
+    #
+    pix_height = np.argmin(np.abs(tau_cube - tau), axis=2) * u.pix - base_height_pix
+    Mm_height = (pix_height * snapshot.ds[2]).to(u.Mm)
+    #
+    nf2_z = nf2_height['coords'][:, :, 0, 2]
+    #
+    vmax = max(Mm_height.max(), nf2_z.max()).value
+    #
+    ax = row[0]
+    im = ax.imshow(Mm_height.value.T, vmin=0, vmax=vmax, cmap='inferno', origin='lower')
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    fig.colorbar(im, cax=cax, label='Height [Mm]')
+    ax.set_title(f'Tau: {tau:.1e}')
+    #
+    ax = row[1]
+    im = ax.imshow(nf2_z.value.T, vmin=0, vmax=vmax, cmap='inferno', origin='lower')
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    fig.colorbar(im, cax=cax, label='Height [Mm]')
+    ax.set_title(f'NF2 Height {i + 1:1d}')
+
+fig.tight_layout()
+fig.savefig(os.path.join(args.output, f'compare_boundary.jpg'))
+plt.close(fig)
\ No newline at end of file
diff --git a/nf2/evaluation/muram/compare_currents.py b/nf2/evaluation/muram/compare_currents.py
new file mode 100644
index 0000000..86e8687
--- /dev/null
+++ b/nf2/evaluation/muram/compare_currents.py
@@ -0,0 +1,216 @@
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.evaluation.metric import curl
+from nf2.evaluation.output import CartesianOutput
+from nf2.loader.muram import MURaMSnapshot
+
+from astropy import constants as const
+from astropy import units as u
+
+muram_path = '/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/3D'
+iteration = 474000
+
+
+muram_snapshot = MURaMSnapshot(muram_path, iteration)
+muram_cube = muram_snapshot.load_cube()
+b_muram = muram_cube['B']
+
+nf2_single_height = CartesianOutput('/glade/work/rjarolim/nf2/multi_height/muram_mfr_single_height_v02/extrapolation_result.nf2')
+nf2_multi_height = CartesianOutput('/glade/work/rjarolim/nf2/multi_height/muram_mfr_v03/extrapolation_result.nf2')
+nf2_disambiguation = CartesianOutput('/glade/work/rjarolim/nf2/multi_height/muram_mfr_disambiguation_v04/extrapolation_result.nf2')
+
+nf2_res = 0.192
+muram_res = 0.192
+res_single_height = nf2_single_height.load_cube(Mm_per_pixel=nf2_res, progress=True, height_range=[0, b_muram.shape[2] * muram_res])
+res_multi_height = nf2_multi_height.load_cube(Mm_per_pixel=nf2_res, progress=True, height_range=[0, b_muram.shape[2] * muram_res])
+res_disambiguation = nf2_disambiguation.load_cube(Mm_per_pixel=nf2_res, progress=True, height_range=[0, b_muram.shape[2] * muram_res])
+
+# integrated currents
+j_muram = curl(b_muram) * u.G / (muram_res * u.Mm) * const.c / (4 * np.pi) # Mm_per_pixel
+j_muram = (j_muram).to(u.G / u.s)
+
+offset_z = int(5 / nf2_res)
+integrated_j_single_height = np.linalg.norm(res_single_height['J'], axis=-1)[:, :, offset_z:].sum(axis=2) * nf2_single_height.Mm_per_pixel
+integrated_j_multi_height = np.linalg.norm(res_multi_height['J'], axis=-1)[:, :, offset_z:].sum(axis=2) * nf2_multi_height.Mm_per_pixel
+integrated_j_disambiguation = np.linalg.norm(res_disambiguation['J'], axis=-1)[:, :, offset_z:].sum(axis=2) * nf2_disambiguation.Mm_per_pixel
+offset_z = int(5 / muram_res)
+integrated_j_muram = np.linalg.norm(j_muram, axis=-1)[:, :, offset_z:].sum(2) * muram_res * u.Mm / u.pix
+
+# plot integrated currents
+norm = LogNorm(vmin=1e1, vmax=1e6)
+
+fig, axs = plt.subplots(3, 4, figsize=(15, 10))
+
+
+ax = axs[0, 0]
+extent = [0, muram_res * integrated_j_muram.shape[0], 0, muram_res * integrated_j_muram.shape[1]]
+im = ax.imshow(integrated_j_muram.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('MURaM')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[0, 1]
+extent = [0, nf2_res * integrated_j_single_height.shape[0], 0, nf2_res * integrated_j_single_height.shape[1]]
+im = ax.imshow(integrated_j_single_height.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Single Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[0, 2]
+extent = [0, nf2_res * integrated_j_multi_height.shape[0], 0, nf2_res * integrated_j_multi_height.shape[1]]
+im = ax.imshow(integrated_j_multi_height.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Multi Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[0, 3]
+extent = [0, nf2_res * integrated_j_disambiguation.shape[0], 0, nf2_res * integrated_j_disambiguation.shape[1]]
+im = ax.imshow(integrated_j_disambiguation.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Auto-Disambiguation')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+x_slice_Mm = 35
+[a.set_xlabel('X [Mm]') for a in axs[0]]
+[a.set_ylabel('Y [Mm]') for a in axs[0]]
+[a.axvline(x_slice_Mm, color='white', linestyle='--') for a in axs[0]]
+
+# plot slices
+x = int(x_slice_Mm / muram_res)
+muram_slice = np.linalg.norm(j_muram, axis=-1)[x]
+muram_b_slice = np.linalg.norm(b_muram[..., :2], axis=-1)[x]
+x = int(x_slice_Mm / nf2_res)
+single_height_slice = np.linalg.norm(res_single_height['J'], axis=-1)[x]
+single_height_b_slice = np.linalg.norm(res_single_height['B'][..., :2], axis=-1)[x]
+multi_height_slice = np.linalg.norm(res_multi_height['J'], axis=-1)[x]
+multi_height_b_slice = np.linalg.norm(res_multi_height['B'][..., :2], axis=-1)[x]
+disambiguation_slice = np.linalg.norm(res_disambiguation['J'], axis=-1)[x]
+disambiguation_b_slice = np.linalg.norm(res_disambiguation['B'][..., :2], axis=-1)[x]
+
+norm = LogNorm(vmin=1e1, vmax=2e4)
+
+ax = axs[1, 0]
+extent = [0, muram_res * muram_slice.shape[0], 0, muram_res * muram_slice.shape[1]]
+im = ax.imshow(muram_slice.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('MURaM')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[1, 1]
+extent = [0, nf2_res * single_height_slice.shape[0], 0, nf2_res * single_height_slice.shape[1]]
+im = ax.imshow(single_height_slice.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Single Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[1, 2]
+extent = [0, nf2_res * multi_height_slice.shape[0], 0, nf2_res * multi_height_slice.shape[1]]
+im = ax.imshow(multi_height_slice.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Multi Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[1, 3]
+extent = [0, nf2_res * disambiguation_slice.shape[0], 0, nf2_res * disambiguation_slice.shape[1]]
+im = ax.imshow(disambiguation_slice.value.T, cmap='plasma', norm=norm, extent=extent, origin='lower')
+ax.set_title('Auto-Disambiguation')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+[a.set_xlabel('X [Mm]') for a in axs[1]]
+[a.set_ylabel('Z [Mm]') for a in axs[1]]
+
+b_norm = LogNorm(vmin=10, vmax=1e3)
+ax = axs[2, 0]
+scaling = (muram_b_slice.shape[1] * muram_res) ** 2
+extent = [0, muram_res * muram_b_slice.shape[0], 0, muram_res * muram_b_slice.shape[1]]
+im = ax.imshow(muram_b_slice.T, cmap='cividis', extent=extent, origin='lower', norm=b_norm)
+ax.set_title('MURaM')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[2, 1]
+scaling = (single_height_b_slice.shape[1] * nf2_res) ** 2
+extent = [0, nf2_res * single_height_b_slice.shape[0], 0, nf2_res * single_height_b_slice.shape[1]]
+im = ax.imshow(single_height_b_slice.value.T, cmap='cividis', extent=extent, origin='lower', norm=b_norm)
+ax.set_title('Single Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[2, 2]
+extent = [0, nf2_res * multi_height_b_slice.shape[0], 0, nf2_res * multi_height_b_slice.shape[1]]
+im = ax.imshow(multi_height_b_slice.value.T, cmap='cividis', extent=extent, origin='lower', norm=b_norm)
+ax.set_title('Multi Height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[2, 3]
+extent = [0, nf2_res * disambiguation_b_slice.shape[0], 0, nf2_res * disambiguation_b_slice.shape[1]]
+im = ax.imshow(disambiguation_b_slice.value.T, cmap='cividis', extent=extent, origin='lower', norm=b_norm)
+ax.set_title('Auto-Disambiguation')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+[a.set_xlabel('X [Mm]') for a in axs[2]]
+[a.set_ylabel('Z [Mm]') for a in axs[2]]
+
+fig.tight_layout()
+fig.savefig('/glade/work/rjarolim/nf2/multi_height/slices.jpg', dpi=300)
+plt.close(fig)
+
+height_muram = np.linspace(0, muram_res * muram_slice.shape[1], muram_slice.shape[1])
+height_nf2 = np.linspace(0, nf2_res * disambiguation_slice.shape[1], disambiguation_slice.shape[1])
+
+profile_muram = muram_slice.mean(axis=0)
+profile_single_height = single_height_slice.mean(axis=0)
+profile_multi_height = multi_height_slice.mean(axis=0)
+profile_disambiguation = disambiguation_slice.mean(axis=0)
+
+profile_bh_muram = muram_b_slice.mean(axis=0)
+profile_bh_single_height = single_height_b_slice.mean(axis=0)
+profile_bh_multi_height = multi_height_b_slice.mean(axis=0)
+profile_bh_disambiguation = disambiguation_b_slice.mean(axis=0)
+
+fig, axs = plt.subplots(1, 2, figsize=(8, 4))
+
+ax = axs[0]
+ax.plot(profile_muram, height_muram, label='MURaM')
+ax.plot(profile_single_height, height_nf2, label='Single Height')
+ax.plot(profile_multi_height, height_nf2, label='Multi Height')
+ax.plot(profile_disambiguation, height_nf2, label='Auto-Disambiguation')
+ax.set_xlabel('Current [G/s]')
+ax.set_ylabel('Z [Mm]')
+ax.semilogx()
+ax.set_ylim(0)
+
+ax = axs[1]
+ax.plot(profile_bh_muram, height_muram, label='MURaM')
+ax.plot(profile_bh_single_height, height_nf2, label='Single Height')
+ax.plot(profile_bh_multi_height, height_nf2, label='Multi Height')
+ax.plot(profile_bh_disambiguation, height_nf2, label='Auto-Disambiguation')
+ax.set_xlabel('B$_{||}$ [G]')
+ax.set_ylabel('Z [Mm]')
+ax.semilogx()
+ax.set_ylim(0)
+
+# add single legend for both plots
+axs[1].legend(bbox_to_anchor=(1.1, 1.05))
+
+fig.tight_layout()
+fig.savefig('/glade/work/rjarolim/nf2/multi_height/profiles.jpg', dpi=300)
+plt.close(fig)
\ No newline at end of file
diff --git a/nf2/evaluation/muram/cube_height.py b/nf2/evaluation/muram/cube_height.py
new file mode 100644
index 0000000..d63a432
--- /dev/null
+++ b/nf2/evaluation/muram/cube_height.py
@@ -0,0 +1,47 @@
+import glob
+
+import numpy as np
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.loader.muram import MURaMSnapshot
+import argparse
+
+from astropy import units as u
+
+parser = argparse.ArgumentParser(description='Evaluate MURaM snapshot.')
+parser.add_argument('--source_path', type=str, help='path to the MURaM simulation.')
+parser.add_argument('--iteration', type=int, help='iteration of the snapshot.')
+args = parser.parse_args()
+
+
+snapshot = MURaMSnapshot(args.source_path, args.iteration)
+tau_cube = snapshot.tau
+
+pix_height = np.argmin(np.abs(tau_cube - 1), axis=2) * u.pix
+Mm_height = (pix_height * snapshot.ds[2]).to(u.Mm)
+base_height_Mm = Mm_height.min()
+base_height_pix = pix_height.min()
+cube_height = (snapshot.shape[2] * u.pix * snapshot.ds[2]).to(u.Mm)
+print(f'Cube Height: {cube_height} [Mm]; {snapshot.shape[2]} [pix]')
+print(f'Base Height: {base_height_Mm} [Mm]; {base_height_pix} [pix]')
+
+fig, axs = plt.subplots(1, 4, figsize=(16, 4))
+
+for ax, tau in zip(axs, [1.0, 0.001, 0.000100, 0.000001]):
+    pix_height = np.argmin(np.abs(tau_cube - tau), axis=2) * u.pix - base_height_pix
+    Mm_height = (pix_height * snapshot.ds[2]).to(u.Mm)
+    #
+    im = ax.imshow(Mm_height.value.T, vmin=0, cmap='inferno', origin='lower')
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    fig.colorbar(im, cax=cax, label='Height [Mm]')
+    ax.set_title(f'Tau: {tau:.1e}')
+    #
+    print(f'Tau: {tau}')
+    print(f'Height: {Mm_height.min()} - {Mm_height.max()}; {Mm_height.mean()}')
+    print(f'Height: {pix_height.min()} - {pix_height.max()}; {pix_height.mean()}')
+
+fig.tight_layout()
+fig.savefig(f'/glade/work/rjarolim/nf2/multi_height/tau_height.jpg')
+plt.close(fig)
\ No newline at end of file
diff --git a/nf2/evaluation/muram/height_metrics.py b/nf2/evaluation/muram/height_metrics.py
index 2463563..8a04bc1 100644
--- a/nf2/evaluation/muram/height_metrics.py
+++ b/nf2/evaluation/muram/height_metrics.py
@@ -35,39 +35,39 @@
 axs[-1].axvline(1, linestyle='--', color='black')
 
 ckpt_paths = [
-    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation/extrapolation_result.nf2',
-    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation_pf/extrapolation_result.nf2',
-    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_fixed/extrapolation_result.nf2',
-    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_ideal/extrapolation_result.nf2',
-    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau/extrapolation_result.nf2',
-    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bxyz_split/extrapolation_result.nf2',
-    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bz_2epochs/extrapolation_result.nf2',
+    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation/extrapolation_result.nf2',
     '/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation_pf/extrapolation_result.nf2',
+    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_fixed/extrapolation_result.nf2',
+    '/gpfs/gpfs0/robert.jarolim/multi_height/muram_ideal/extrapolation_result.nf2',
+    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau/extrapolation_result.nf2',
+    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bxyz_split/extrapolation_result.nf2',
+    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bz_2epochs/extrapolation_result.nf2',
+    # '/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation_pf/extrapolation_result.nf2',
 ]
 labels = [
-    # 'Extrapolation',
-    # 'Extrapolation - PB',
-    # 'Fixed Heights',
-    # 'Mapped Heights',
-    r'Realistic vector',
-    r'Realistic split',
-    r'Realistic LOS',
+    'Extrapolation',
     'Extrapolation - PB',
+    'Fixed Heights',
+    'Mapped Heights',
+    # r'Realistic vector',
+    # r'Realistic split',
+    # r'Realistic LOS',
+    # 'Extrapolation - PB',
 ]
 
 linestyle = [
-            # ':', ':',
-             # '--', '--',
-             'dashdot', 'dashdot',
-             'dashdot',
-             ':',
+            ':', ':',
+             '--', '--',
+             # 'dashdot', 'dashdot',
+             # 'dashdot',
+             # ':',
              ]
 
 # assert same length
 assert len(ckpt_paths) == len(labels) == len(linestyle)
 
-# colors = [f'C{i + 1}' for i in range(len(ckpt_paths))]
-colors = ['C5', 'C6', 'C7', 'C2']
+colors = ['C1', 'C2', 'C3', 'C4']
+# colors = ['C5', 'C6', 'C7', 'C2']
 
 def _plot(b, label, c, ls):
     c_vec = np.sum((B * b).sum(-1), (0, 1)) / np.sqrt((B ** 2).sum(-1).sum((0, 1)) * (b ** 2).sum(-1).sum((0, 1)))
@@ -112,7 +112,7 @@ def _plot(b, label, c, ls):
 axs[0].set_xlim(0.7, 1)
 axs[1].set_xlim(0.25, .85)
 axs[2].set_xlim(0., .8)
-axs[3].set_xlim(-0.5, .5)
+axs[3].set_xlim(-0.5, .6)
 axs[4].set_xlim(0.5, 1.5)
 
 [ax.set_ylim(0, heights.max()) for ax in axs]
@@ -120,5 +120,5 @@ def _plot(b, label, c, ls):
 
 # fig.tight_layout()
 fig.subplots_adjust(bottom=.27)
-fig.savefig(f'{base_path}/comparison_metrics.jpg', dpi=300)
+fig.savefig(f'{base_path}/comparison_metrics_full_v2.jpg', dpi=300)
 plt.close(fig)
diff --git a/nf2/evaluation/muram/muram_to_vtk.py b/nf2/evaluation/muram/muram_to_vtk.py
new file mode 100644
index 0000000..05e9d91
--- /dev/null
+++ b/nf2/evaluation/muram/muram_to_vtk.py
@@ -0,0 +1,33 @@
+import argparse
+
+import numpy as np
+from astropy.nddata import block_reduce
+
+from nf2.evaluation.metric import curl
+from nf2.evaluation.output import CartesianOutput
+from nf2.evaluation.vtk import save_vtk
+from nf2.loader.muram import MURaMSnapshot
+
+from astropy import units as u
+
+from astropy import constants as const
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--source_path', type=str, help='path to the MURaM simulation.')
+parser.add_argument('--iteration', type=int, help='iteration of the snapshot.')
+parser.add_argument('--vtk_path', type=str, help='path to the target VTK file', required=False, default=None)
+parser.add_argument('--Mm_per_pixel', type=float, help='spatial resolution (needs to be a multiple of 0.192)', required=False, default=0.192 * 2)
+args = parser.parse_args()
+
+Mm_per_pixel = args.Mm_per_pixel * u.Mm / u.pix
+
+snapshot = MURaMSnapshot(args.source_path, args.iteration)
+muram_cube = snapshot.load_base(args.Mm_per_pixel * u.Mm / u.pix, base_height=0)
+b = muram_cube['B']
+j = curl(b) * u.G / (args.Mm_per_pixel * u.Mm) * const.c / (4 * np.pi) # Mm_per_pixel
+j = j.to(u.G / u.s)
+
+azimuth = np.arctan2(b[..., 1], b[..., 0])
+
+save_vtk(args.vtk_path, vectors={'B': b, 'J': j.value},
+         scalars={'tau': muram_cube['tau'], 'azimuth': azimuth}, Mm_per_pix=args.Mm_per_pixel)
diff --git a/nf2/evaluation/nf2_to_fits.py b/nf2/evaluation/nf2_to_fits.py
deleted file mode 100644
index 1addd6b..0000000
--- a/nf2/evaluation/nf2_to_fits.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import argparse
-
-import torch
-from astropy.io import fits
-
-from nf2.evaluation.unpack import load_cube
-
-parser = argparse.ArgumentParser(description='Convert NF2 file to fits.')
-parser.add_argument('nf2_path', type=str, help='path to the source NF2 file')
-parser.add_argument('fits_path', type=str, help='path to the target numpy file')
-parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-
-args = parser.parse_args()
-nf2_path = args.nf2_path
-strides = args.strides
-fits_path = args.fits_path
-
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-
-b = load_cube(nf2_path, device, progress=True, strides=strides)
-
-hdu = fits.PrimaryHDU(b)
-hdul = fits.HDUList([hdu])
-hdul.writeto(fits_path)
-
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
diff --git a/nf2/evaluation/nf2_to_npy.py b/nf2/evaluation/nf2_to_npy.py
deleted file mode 100644
index 6c651aa..0000000
--- a/nf2/evaluation/nf2_to_npy.py
+++ /dev/null
@@ -1,26 +0,0 @@
-import argparse
-
-import numpy as np
-import torch
-
-from nf2.evaluation.unpack import load_cube
-
-parser = argparse.ArgumentParser(description='Convert NF2 file to npy.')
-parser.add_argument('nf2_path', type=str, help='path to the source NF2 file')
-parser.add_argument('npy_path', type=str, help='path to the target numpy file')
-parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-
-args = parser.parse_args()
-nf2_path = args.nf2_path
-strides = args.strides
-npy_path = args.npy_path
-
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-
-b = load_cube(nf2_path, device, progress=True, strides=strides)
-
-np.save(npy_path, b)
-
-
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
diff --git a/nf2/evaluation/nf2_to_vtk.py b/nf2/evaluation/nf2_to_vtk.py
deleted file mode 100644
index b221c34..0000000
--- a/nf2/evaluation/nf2_to_vtk.py
+++ /dev/null
@@ -1,32 +0,0 @@
-import argparse
-
-import torch
-
-from nf2.evaluation.unpack import load_cube, load_height_cube
-from nf2.evaluation.vtk import save_vtk
-
-parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
-parser.add_argument('nf2_path', type=str, help='path to the source NF2 file')
-parser.add_argument('vtk_path', type=str, help='path to the target VTK file')
-parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-
-args = parser.parse_args()
-nf2_path = args.nf2_path
-strides = args.strides
-vtk_path = args.vtk_path
-
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
-
-b = load_cube(nf2_path, device, progress=True, strides=strides)
-
-state = torch.load(nf2_path, map_location=device)
-if state['height_mapping_model'] is not None:
-    tau = load_height_cube(nf2_path, device=device, progress=True, strides=strides)
-else:
-    tau = None
-
-save_vtk(b, vtk_path, 'B', scalar=tau, scalar_name='tau',
-         Mm_per_pix=state['Mm_per_pixel'] * strides)
-
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
diff --git a/nf2/evaluation/output.py b/nf2/evaluation/output.py
new file mode 100644
index 0000000..991b991
--- /dev/null
+++ b/nf2/evaluation/output.py
@@ -0,0 +1,489 @@
+import numpy as np
+import torch
+from astropy import units as u, constants
+from astropy.coordinates import SkyCoord
+from dateutil.parser import parse
+from sunpy.coordinates import frames
+from sunpy.map import Map
+from torch import nn
+from tqdm import tqdm
+
+from nf2.data.util import spherical_to_cartesian, cartesian_to_spherical, vector_cartesian_to_spherical
+from nf2.evaluation.energy import get_free_mag_energy
+from nf2.evaluation.metric import energy, normalized_divergence
+from nf2.train.model import VectorPotentialModel, calculate_current_from_jacobian
+
+
+def current_density(jac_matrix, **kwargs):
+    j = calculate_current_from_jacobian(jac_matrix, f=np) * constants.c / (4 * np.pi)
+    return j.to(u.G / u.s)
+
+
+def b_nabla_bz(b, jac_matrix, **kwargs):
+    # compute B * nabla * Bz
+    bx = b[..., 0]
+    by = b[..., 1]
+    bz = b[..., 2]
+    dBx_dx = jac_matrix[..., 0, 0]
+    dBx_dy = jac_matrix[..., 0, 1]
+    dBx_dz = jac_matrix[..., 0, 2]
+    dBy_dx = jac_matrix[..., 1, 0]
+    dBy_dy = jac_matrix[..., 1, 1]
+    dBy_dz = jac_matrix[..., 1, 2]
+    dBz_dx = jac_matrix[..., 2, 0]
+    dBz_dy = jac_matrix[..., 2, 1]
+    dBz_dz = jac_matrix[..., 2, 2]
+    #
+    norm_B = np.linalg.norm(b, axis=-1) + 1e-10 * b.unit
+
+    dnormB_dx = - norm_B ** -3 * (bx * dBx_dx + by * dBy_dx + bz * dBz_dx)
+    dnormB_dy = - norm_B ** -3 * (bx * dBx_dy + by * dBy_dy + bz * dBz_dy)
+    dnormB_dz = - norm_B ** -3 * (bx * dBx_dz + by * dBy_dz + bz * dBz_dz)
+
+    b_nabla_bz = (bx * dBz_dx + by * dBz_dy + bz * dBz_dz) / norm_B ** 2 + \
+                 (bz / norm_B) * (bx * dnormB_dx + by * dnormB_dy + bz * dnormB_dz)
+    b_nabla_bz = b_nabla_bz.to(1 / u.Mm)
+    return b_nabla_bz
+
+def magnetic_helicity(b, a, **kwargs):
+    return np.sum(a * b, axis=-1)
+
+def twist(b, jac_matrix, **kwargs):
+    j = calculate_current_from_jacobian(jac_matrix, f=np)
+    twist = np.linalg.norm(j, axis=-1) / np.linalg.norm(b, axis=-1)
+    threshold = np.linalg.norm(b, axis=-1).to_value(u.G) < 10 # set twist to 0 for weak fields (coronal holes)
+    twist[threshold] = 0 * u.Mm ** -1
+    twist = twist.to(u.Mm ** -1)
+    return twist
+
+
+def spherical_energy_gradient(b, jac_matrix, coords, **kwargs):
+    dBx_dx = jac_matrix[..., 0, 0]
+    dBy_dx = jac_matrix[..., 1, 0]
+    dBz_dx = jac_matrix[..., 2, 0]
+    dBx_dy = jac_matrix[..., 0, 1]
+    dBy_dy = jac_matrix[..., 1, 1]
+    dBz_dy = jac_matrix[..., 2, 1]
+    dBx_dz = jac_matrix[..., 0, 2]
+    dBy_dz = jac_matrix[..., 1, 2]
+    dBz_dz = jac_matrix[..., 2, 2]
+    # E = b^2 = b_x^2 + b_y^2 + b_z^2
+    # dE/dx = 2 * (b_x * dBx_dx + b_y * dBy_dx + b_z * dBz_dx)
+    # dE/dy = 2 * (b_x * dBx_dy + b_y * dBy_dy + b_z * dBz_dy)
+    # dE/dz = 2 * (b_x * dBx_dz + b_y * dBy_dz + b_z * dBz_dz)
+    dE_dx = 2 * (b[..., 0] * dBx_dx + b[..., 1] * dBy_dx + b[..., 2] * dBz_dx)
+    dE_dy = 2 * (b[..., 0] * dBx_dy + b[..., 1] * dBy_dy + b[..., 2] * dBz_dy)
+    dE_dz = 2 * (b[..., 0] * dBx_dz + b[..., 1] * dBy_dz + b[..., 2] * dBz_dz)
+
+    coords_spherical = cartesian_to_spherical(coords, f=np)
+    t = coords_spherical[..., 1]
+    p = coords_spherical[..., 2]
+    dE_dr = (np.sin(t) * np.cos(p)) * dE_dx + \
+            (np.sin(t) * np.sin(p)) * dE_dy + \
+            np.cos(p) * dE_dz
+
+    return dE_dr
+
+
+def energy_gradient(b, jac_matrix, **kwargs):
+    dBx_dz = jac_matrix[..., 0, 2]
+    dBy_dz = jac_matrix[..., 1, 2]
+    dBz_dz = jac_matrix[..., 2, 2]
+    # E = b^2 = b_x^2 + b_y^2 + b_z^2
+    # dE/dz = 2 * (b_x * dBx_dz + b_y * dBy_dz + b_z * dBz_dz)
+    dE_dz = 2 * (b[..., 0] * dBx_dz + b[..., 1] * dBy_dz + b[..., 2] * dBz_dz)
+
+    return dE_dz
+
+
+def los_trv_azi(b, **kwargs):
+    bx, by, bz = b[..., 0], b[..., 1], b[..., 2]
+    b_los = bz.to_value(u.G)
+    b_trv = np.sqrt(bx ** 2 + by ** 2).to_value(u.G)
+    azimuth = np.arctan2(by, bx).to_value(u.deg)
+    return np.stack([b_los, b_trv, azimuth], -1)
+
+def free_energy(b, **kwargs):
+    free_energy = get_free_mag_energy(b.to_value(u.G)) * u.erg * u.cm ** -3
+    return free_energy
+
+class BaseOutput:
+
+    def __init__(self, checkpoint, device=None):
+        if device is None:
+            device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+        self.state = torch.load(checkpoint, map_location=device)
+        model = self.state['model']
+        self._requires_grad = isinstance(model, VectorPotentialModel)
+        self.model = nn.DataParallel(model) if torch.cuda.device_count() > 1 else model
+        self.spatial_norm = 1
+        self.device = device
+        self.c = constants.c
+
+    @property
+    def G_per_dB(self):
+        return self.state['data']['G_per_dB'] * u.G
+
+    @property
+    def m_per_ds(self):
+        return (self.state['data']['Mm_per_ds'] * u.Mm).to(u.m)
+
+    def load_coords(self, coords, batch_size=int(2 ** 12), progress=False, compute_jacobian=True,
+                    metrics={'j': current_density}):
+        batch_size = batch_size * torch.cuda.device_count() if torch.cuda.is_available() else batch_size
+        def _load(coords):
+            # normalize and to tensor
+            coords = torch.tensor(coords / self.spatial_norm, dtype=torch.float32)
+            coords_shape = coords.shape
+            coords = coords.reshape((-1, 3))
+
+            model_out = {}
+            it = range(int(np.ceil(coords.shape[0] / batch_size)))
+            it = tqdm(it, desc='Load NF2') if progress else it
+            for k in it:
+                self.model.zero_grad()
+                coord = coords[k * batch_size: (k + 1) * batch_size]
+                coord = coord.to(self.device)
+                coord.requires_grad = True
+                result = self.model(coord, compute_jacobian=compute_jacobian)
+                for k, v in result.items():
+                    if k not in model_out:
+                        model_out[k] = []
+                    model_out[k] += [v.detach().cpu()]
+
+            model_out = {k: torch.cat(v) for k, v in model_out.items()}
+            model_out = {k: v.reshape(*coords_shape[:-1], *v.shape[1:]).numpy() for k, v in model_out.items()}
+
+            model_out['b'] = model_out['b'] * self.G_per_dB
+            if 'a' in model_out:
+                model_out['a'] = model_out['a'] * self.G_per_dB * self.m_per_ds
+            return model_out
+
+        if compute_jacobian or self._requires_grad:
+            model_out = _load(coords)
+            jac_matrix = model_out['jac_matrix']
+            jac_matrix = jac_matrix * self.G_per_dB / self.m_per_ds
+            model_out['jac_matrix'] = jac_matrix
+
+            state = {**model_out, 'coords': coords}
+
+            for k, f in metrics.items():
+                model_out[k] = f(**state)
+
+            return model_out
+        else:
+            with torch.no_grad():
+                return _load(coords)
+
+
+class CartesianOutput(BaseOutput):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.state['data']['type'] == 'cartesian', 'Requires cartesian NF2 data!'
+
+        self.coord_range = self.state['data']['coord_range']
+        self.coord_range = self.coord_range[0] if isinstance(self.coord_range, list) else self.coord_range
+        self.max_height = self.state['data']['max_height']
+        self.ds_per_pixel = self.state['data']['ds_per_pixel']
+        self.ds_per_pixel = self.ds_per_pixel[0] if isinstance(self.ds_per_pixel, list) else self.ds_per_pixel
+        self.Mm_per_ds = self.state['data']['Mm_per_ds']
+        self.Mm_per_pixel = self.ds_per_pixel * self.Mm_per_ds
+        self.wcs = self.state['data']['wcs'] if 'wcs' in self.state['data'] else None
+        self.time = parse(self.wcs[0].wcs.dateobs) if (self.wcs is not None) and len(self.wcs) >= 1 else None
+        self.data_config = self.state['data']
+
+    def load_cube(self, height_range=None, Mm_per_pixel=None, **kwargs):
+        x_min, x_max = self.coord_range[0]
+        y_min, y_max = self.coord_range[1]
+        z_min, z_max = (0, self.max_height / self.Mm_per_ds) if height_range is None else (h / self.Mm_per_ds for h in
+                                                                                           height_range)
+
+        Mm_per_pixel = self.Mm_per_pixel if Mm_per_pixel is None else Mm_per_pixel
+        pixel_per_ds = self.Mm_per_ds / Mm_per_pixel
+
+        coords = np.stack(
+            np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixel_per_ds) + 1),
+                        np.linspace(y_min, y_max, int((y_max - y_min) * pixel_per_ds) + 1),
+                        np.linspace(z_min, z_max, int((z_max - z_min) * pixel_per_ds) + 1), indexing='ij'), -1)
+
+        model_out = self.load_coords(coords, **kwargs)
+
+        return {**model_out, 'coords': coords, 'Mm_per_pixel': Mm_per_pixel}
+
+    def load_boundary(self, Mm_per_pixel=None, **kwargs):
+        x_min, x_max = self.coord_range[0]
+        y_min, y_max = self.coord_range[1]
+
+        Mm_per_pixel = self.Mm_per_pixel if Mm_per_pixel is None else Mm_per_pixel
+        pixel_per_ds = self.Mm_per_ds / Mm_per_pixel
+
+        coords = np.stack(
+            np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixel_per_ds + 1)),
+                        np.linspace(y_min, y_max, int((y_max - y_min) * pixel_per_ds + 1)),
+                        np.zeros((1,), dtype=np.float32), indexing='ij'), -1)
+        coords = coords[:, :, 0]
+
+        model_out = self.load_coords(coords, **kwargs)
+
+        return {**model_out, 'coords': coords, 'Mm_per_pixel': Mm_per_pixel}
+
+    def load_maps(self, **kwargs):
+        model_out = self.load_cube(**kwargs)
+
+        j_map = np.linalg.norm(model_out['j'], axis=-1).sum(axis=-1)
+        b_map = np.linalg.norm(model_out['b'], axis=-1).sum(axis=-1)
+        energy_map = energy(model_out['b']).sum(axis=-1)
+        free_energy_map = get_free_mag_energy(model_out['b']).sum(axis=-1)
+
+        return {'b': Map(b_map, wcs=self.wcs),
+                'j': Map(j_map, wcs=self.wcs),
+                'energy': Map(energy_map, wcs=self.wcs),
+                'free_energy': Map(free_energy_map, wcs=self.wcs)}
+
+    def trace_bottom(self, Mm_per_pixel=None, **kwargs):
+        x_min, x_max = self.coord_range[0]
+        y_min, y_max = self.coord_range[1]
+
+        Mm_per_pixel = self.Mm_per_pixel if Mm_per_pixel is None else Mm_per_pixel
+        pixel_per_ds = self.Mm_per_ds / Mm_per_pixel
+
+        coords = np.stack(
+            np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixel_per_ds + 1)),
+                        np.linspace(y_min, y_max, int((y_max - y_min) * pixel_per_ds + 1)),
+                        np.zeros((1,), dtype=np.float32), indexing='ij'), -1)
+        forward_trace = self.trace(coords, **kwargs)
+        backward_trace = self.trace(coords, direction=-1, **kwargs)
+        traces = {i: list(reversed(backward_trace[i][1:])) + forward_trace[i] for i in forward_trace.keys()}
+        return traces
+
+    def trace(self, start_coords, direction=1, max_iterations=None, **kwargs):
+        '''
+        Trace the field line from the given coordinates using a 4th order Runge-Kutta method.
+        '''
+        base_step = np.array([self.ds_per_pixel / 4], dtype=np.float32).reshape((1, 1))  # quarter of a pixel
+
+        x_min, x_max = self.coord_range[0]
+        y_min, y_max = self.coord_range[1]
+        z_min, z_max = (0, self.max_height / self.Mm_per_ds)
+
+        max_iterations = ((x_max - x_min) + (y_max - y_min) + (
+                    z_max - z_min)) / base_step * 3 if max_iterations is None else max_iterations
+
+        field_lines = {i: [c] for i, c in enumerate(start_coords)}
+
+        coords = start_coords
+        iteration = 0
+        while np.isnan(coords).sum() != 0 and iteration < max_iterations:
+            nan_mask = np.isnan(coords)  # only trace non-nan coordinates
+
+            coords_k1 = coords[~nan_mask]
+            model_out_k1 = self.load_coords(coords_k1, **kwargs)
+            b_k1 = model_out_k1['b']
+            k1 = b_k1 / np.linalg.norm(b_k1, axis=-1, keepdims=True)
+
+            coords_k2 = coords_k1 + k1 * base_step / 2
+            # check this too
+            model_out_k2 = self.load_coords(coords_k2, **kwargs)
+            b_k2 = model_out_k2['b'] * np.sign(direction)
+            k2 = b_k2 / np.linalg.norm(b_k2, axis=-1, keepdims=True)
+
+            coords_k3 = coords_k1 + k2 * base_step / 2
+            model_out_k3 = self.load_coords(coords_k3, **kwargs)
+            b_k3 = model_out_k3['b'] * np.sign(direction)
+            k3 = b_k3 / np.linalg.norm(b_k3, axis=-1, keepdims=True)
+
+            coords_k4 = coords_k1 + k3 * base_step
+            model_out_k4 = self.load_coords(coords_k4, **kwargs)
+            b_k4 = model_out_k4['b'] * np.sign(direction)
+            k4 = b_k4 / np.linalg.norm(b_k4, axis=-1, keepdims=True)
+
+            next_coords = coords_k1 + (k1 + 2 * k2 + 2 * k3 + k4) * base_step / 6
+
+            mask = (next_coords[..., 0] >= x_min) & (next_coords[..., 0] <= x_max) & \
+                   (next_coords[..., 1] >= y_min) & (next_coords[..., 1] <= y_max) & \
+                   (next_coords[..., 2] >= z_min) & (next_coords[..., 2] <= z_max)
+            next_coords[~mask] = None
+            # todo write final value at the boundary
+
+            coords[~nan_mask] = next_coords
+
+            for i, c in enumerate(coords):
+                if c is None:
+                    continue
+                field_lines[i].append(c)
+            iteration += 1
+        return field_lines
+
+
+class HeightTransformOutput(CartesianOutput):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert 'transform_module' in self.state, 'Requires transform module!'
+        self.transform_module = self.state['transform_module']
+
+        self.coord_range_list = self.state['data']['coord_range']
+        self.height_mapping_list = self.state['data']['height_mapping']
+        self.ds_per_pixel_list = self.state['data']['ds_per_pixel']
+
+    def load_height_mapping(self, **kwargs):
+        for coord_range, height_mapping, ds_per_pixel in zip(self.coord_range_list, self.height_mapping_list,
+                                                             self.ds_per_pixel_list):
+            x_min, x_max = coord_range[0]
+            y_min, y_max = coord_range[1]
+            z = height_mapping['z'] / self.Mm_per_ds
+
+            pixel_per_ds = 1 / ds_per_pixel
+            coords = np.stack(
+                np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixel_per_ds)),
+                            np.linspace(y_min, y_max, int((y_max - y_min) * pixel_per_ds)),
+                            z, indexing='ij'), -1)
+            height_range = np.zeros((*coords.shape[:-1], 2), dtype=np.float32)
+            height_range[..., 0] = height_mapping['z_min'] / self.Mm_per_ds
+            height_range[..., 1] = height_mapping['z_max'] / self.Mm_per_ds
+
+            model_out = self.load_transformed_coords(coords, height_range, **kwargs)
+            yield {'coords': model_out['coords'] * self.Mm_per_ds * u.Mm,
+                   'original_coords': coords * self.Mm_per_ds * u.Mm,
+                   'height_range': height_range * self.Mm_per_ds * u.Mm}
+
+    def load_transformed_coords(self, coords, height_range, batch_size=int(2 ** 12), progress=False):
+        def _load(coords, height_range):
+            # normalize and to tensor
+            coords = torch.tensor(coords / self.spatial_norm, dtype=torch.float32)
+            coords_shape = coords.shape
+            coords = coords.reshape((-1, 3))
+
+            height_range = torch.tensor(height_range, dtype=torch.float32)
+            height_range = height_range.reshape((-1, 2))
+
+            cube = []
+            it = range(int(np.ceil(coords.shape[0] / batch_size)))
+            it = tqdm(it) if progress else it
+            for k in it:
+                self.transform_module.zero_grad()
+                coord = coords[k * batch_size: (k + 1) * batch_size]
+                coord = coord.to(self.device)
+
+                height_range_batch = height_range[k * batch_size: (k + 1) * batch_size]
+                height_range_batch = height_range_batch.to(self.device)
+
+                transformed_coords = self.transform_module.transform_coords(coord, height_range_batch)
+
+                cube += [transformed_coords.detach().cpu()]
+
+            cube = torch.cat(cube)
+            cube = cube.reshape(*coords_shape).numpy()
+
+            return {'coords': cube}
+
+        with torch.no_grad():
+            return _load(coords, height_range)
+
+
+class SphericalOutput(BaseOutput):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert self.state['data']['type'] == 'spherical', 'Requires spherical NF2 data!'
+
+        self.radius_range = self.state['data']['radius_range']
+        self.spatial_norm = (self.m_per_ds / (1 * u.solRad)).to_value(u.dimensionless_unscaled)
+
+    def load_spherical(self, radius_range: u.Quantity = None,
+                       latitude_range: u.Quantity = (0, np.pi) * u.rad,
+                       longitude_range: u.Quantity = (0, 2 * np.pi) * u.rad,
+                       sampling=[100, 180, 360], **kwargs):
+        radius_range = radius_range if radius_range is not None else self.radius_range
+        spherical_coords = np.stack(
+            np.meshgrid(
+                np.linspace(radius_range[0].to_value(u.solRad), radius_range[1].to_value(u.solRad), sampling[0]),
+                np.linspace(latitude_range[0].to_value(u.rad), latitude_range[1].to_value(u.rad), sampling[1]),
+                np.linspace(longitude_range[0].to_value(u.rad), longitude_range[1].to_value(u.rad), sampling[2]),
+                indexing='ij'), -1)
+        cartesian_coords = spherical_to_cartesian(spherical_coords)
+
+        model_out = self.load_coords(cartesian_coords, **kwargs)
+        return {**model_out, 'coords': cartesian_coords, 'spherical_coords': spherical_coords}
+
+    def load(self,
+             radius_range: u.Quantity = None,
+             latitude_range: u.Quantity = (0, np.pi) * u.rad,
+             longitude_range: u.Quantity = (0, 2 * np.pi),
+             resolution: u.Quantity = 64 * u.pix / u.solRad, nan_value=0, **kwargs):
+        radius_range = radius_range if radius_range is not None else self.radius_range
+        spherical_bounds = np.stack(
+            np.meshgrid(np.linspace(radius_range[0].to_value(u.solRad), radius_range[1].to_value(u.solRad), 50),
+                        np.linspace(latitude_range[0].to_value(u.rad), latitude_range[1].to_value(u.rad), 50),
+                        np.linspace(longitude_range[0].to_value(u.rad), longitude_range[1].to_value(u.rad), 50),
+                        indexing='ij'), -1)
+
+        cartesian_bounds = spherical_to_cartesian(spherical_bounds)
+        x_min, x_max = cartesian_bounds[..., 0].min(), cartesian_bounds[..., 0].max()
+        y_min, y_max = cartesian_bounds[..., 1].min(), cartesian_bounds[..., 1].max()
+        z_min, z_max = cartesian_bounds[..., 2].min(), cartesian_bounds[..., 2].max()
+
+        res = resolution.to_value(u.pix / u.solRad)
+        coords = np.stack(
+            np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * res)),
+                        np.linspace(y_min, y_max, int((y_max - y_min) * res)),
+                        np.linspace(z_min, z_max, int((z_max - z_min) * res)), indexing='ij'), -1)
+        # flipped z axis
+        spherical_coords = cartesian_to_spherical(coords)
+        lat_coord = (spherical_coords[..., 1] % np.pi)
+        lon_coord = (spherical_coords[..., 2] % (2 * np.pi))
+        rad_coord = spherical_coords[..., 0]
+
+        min_lat, max_lat = latitude_range[0].to_value(u.rad), latitude_range[1].to_value(u.rad)
+        min_lon, max_lon = (longitude_range[0].to_value(u.rad), longitude_range[1].to_value(u.rad))
+
+        # only evaluate coordinates in simulation volume
+        if min_lat == max_lat:
+            lat_cond = np.ones_like(lat_coord, dtype=bool)
+        else:
+            lat_cond = (lat_coord >= min_lat) & (lat_coord < max_lat)
+        if min_lon == max_lon:
+            lon_cond = np.ones_like(lon_coord, dtype=bool)
+        else:
+            lon_cond = (lon_coord >= min_lon) & (lon_coord < max_lon)
+            if max_lon > 2 * np.pi:
+                lon_cond = lon_cond | ((lon_coord < max_lon - 2 * np.pi) & (lon_coord >= 0))
+        rad_cond = (rad_coord >= radius_range[0].to_value(u.solRad)) & (rad_coord < radius_range[1].to_value(u.solRad))
+        condition = rad_cond & lat_cond & lon_cond
+        sub_coords = coords[condition]
+
+        cube_shape = coords.shape[:-1]
+        model_out = self.load_coords(sub_coords, **kwargs)
+
+        spherical_out = {'spherical_coords': spherical_coords, 'coords': coords}
+        for k, sub_v in model_out.items():
+            volume = np.ones(cube_shape + sub_v.shape[1:]) * nan_value
+            if hasattr(sub_v, 'unit'): # preserve units
+                volume = volume * sub_v.unit
+            volume[condition] = sub_v
+            spherical_out[k] = volume
+
+        spherical_out['b_rtp'] = vector_cartesian_to_spherical(spherical_out['b'], spherical_coords)
+
+        return spherical_out
+
+    def load_spherical_coords(self, spherical_coords: SkyCoord, **kwargs):
+        cartesian_coords, spherical_coords = self._skycoords_to_cartesian(spherical_coords)
+        model_out = self.load_coords(cartesian_coords, **kwargs)
+        model_out['spherical_coords'] = spherical_coords
+        model_out['coords'] = cartesian_coords
+        return model_out
+
+    def _skycoords_to_cartesian(self, spherical_coords):
+        spherical_coords = spherical_coords.transform_to(frames.HeliographicCarrington)
+        r = spherical_coords.radius
+        r = r * u.solRad if r.unit == u.dimensionless_unscaled else r
+        spherical_coords = np.stack([
+            r.to(u.solRad).value,
+            np.pi / 2 - spherical_coords.lat.to(u.rad).value,
+            spherical_coords.lon.to(u.rad).value,
+        ], -1)
+        cartesian_coords = spherical_to_cartesian(spherical_coords)
+        return cartesian_coords, spherical_coords
diff --git a/nf2/evaluation/plot_carrington_map.py b/nf2/evaluation/plot_carrington_map.py
new file mode 100644
index 0000000..24eebdb
--- /dev/null
+++ b/nf2/evaluation/plot_carrington_map.py
@@ -0,0 +1,48 @@
+import argparse
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from torch import nn
+
+from nf2.data.util import vector_cartesian_to_spherical, spherical_to_cartesian
+from nf2.evaluation.unpack import load_coords
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('nf2_path', type=str, help='path to the source NF2 file')
+parser.add_argument('vtk_path', type=str, help='path to the target VTK file')
+parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
+
+args = parser.parse_args()
+nf2_path = args.nf2_path
+strides = args.strides
+
+nf2_path = '/Users/robert/PycharmProjects/NF2/results/extrapolation_result.nf2'
+height = 1.5
+resolution = 256
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+state = torch.load(nf2_path, map_location=device)
+model = nn.DataParallel(state['model'])
+
+spherical_coords = np.stack(
+    np.meshgrid(np.linspace(1, 1, 1),
+                np.linspace(0 * np.pi, 1 * np.pi, resolution),
+                np.linspace(0 * np.pi, 2 * np.pi, resolution * 2), indexing='ij'), -1)
+coords = spherical_to_cartesian(spherical_coords)
+
+spatial_norm = 1
+cube_shape = coords.shape[:-1]
+b = load_coords(model, spatial_norm, state['b_norm'], coords, device, progress=True)
+
+b_spherical = vector_cartesian_to_spherical(b, spherical_coords)
+
+plt.figure(figsize=(10, 5))
+plt.imshow(b_spherical[0, :, :, 0], cmap='gray', origin='lower', extent=[0, 2, 0, 1])
+# plot grid lines
+plt.grid()
+plt.xlabel('Longitude [rad]')
+plt.ylabel('Latitude [rad]')
+plt.tight_layout()
+plt.show()
\ No newline at end of file
diff --git a/nf2/evaluation/q_map.py b/nf2/evaluation/q_map.py
new file mode 100644
index 0000000..8464655
--- /dev/null
+++ b/nf2/evaluation/q_map.py
@@ -0,0 +1,249 @@
+import argparse
+import os.path
+
+import cupy
+import numpy as np
+from fastqslpy import FastQSL
+from fastqslpy.kernels import compileTraceBlineAdaptive
+from tqdm import tqdm
+
+from nf2.evaluation.output import CartesianOutput
+from nf2.evaluation.vtk import save_vtk
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+    parser.add_argument('--result_path', type=str, help='path to the target VTK file', required=False, default=None)
+    parser.add_argument('--height_range', type=float, nargs=2, help='height range in Mm', required=False, default=None)
+
+    args = parser.parse_args()
+    nf2_path = args.nf2_path
+
+    result_path = args.result_path
+    height_range = args.height_range
+    os.makedirs(result_path, exist_ok=True)
+
+    nf2_out = CartesianOutput(nf2_path)
+    output = nf2_out.load_cube(Mm_per_pixel=0.72, height_range=height_range, progress=True)
+
+    # convert B to cuda array
+    Bx = output["b"][..., 0].astype(np.float32)
+    By = output["b"][..., 1].astype(np.float32)
+    Bz = output["b"][..., 2].astype(np.float32)
+
+    # convert to fortran order
+    Bx_gpu = cupy.array(Bx)
+    By_gpu = cupy.array(By)
+    Bz_gpu = cupy.array(Bz)
+
+    # convert to cuda array
+    Bx_gpu = cupy.asfortranarray(Bx_gpu)
+    By_gpu = cupy.asfortranarray(By_gpu)
+    Bz_gpu = cupy.asfortranarray(Bz_gpu)
+
+    # convert J to cuda array
+    Jx = output["j"][..., 0].astype(np.float32)
+    Jy = output["j"][..., 1].astype(np.float32)
+    Jz = output["j"][..., 2].astype(np.float32)
+
+    # convert to fortran order
+    Jx_gpu = cupy.array(Jx)
+    Jy_gpu = cupy.array(Jy)
+    Jz_gpu = cupy.array(Jz)
+
+    # convert to cuda array
+    Jx_gpu = cupy.asfortranarray(Jx_gpu)
+    Jy_gpu = cupy.asfortranarray(Jy_gpu)
+    Jz_gpu = cupy.asfortranarray(Jz_gpu)
+
+    trace_all_b_line = compileTraceBlineAdaptive()
+
+    # prepare variables
+    BshapeN = np.zeros(3, dtype=np.int32)
+    BshapeN[:] = Bx.shape
+    print(BshapeN)
+    BshapeN = cupy.array(BshapeN)
+
+    interp_ratio = 3
+    stride_step = 1 / interp_ratio
+    x_range = [0, Bx.shape[0]]
+    y_range = [0, Bx.shape[1]]
+    z_range = [0, Bx.shape[2]]
+    # z_range = [0,1]
+
+    x_i = cupy.linspace(*x_range, np.uint(interp_ratio * (x_range[1] - x_range[0])), dtype=cupy.float32)
+    y_i = cupy.linspace(*y_range, np.uint(interp_ratio * (y_range[1] - y_range[0])), dtype=cupy.float32)
+    z_i = cupy.linspace(*z_range, np.uint(interp_ratio * (z_range[1] - z_range[0])), dtype=cupy.float32)
+
+    flag_twist = cupy.array([True], dtype=cupy.bool_)
+    flag_twist_false = cupy.array([False], dtype=cupy.bool_)
+    dummy = cupy.zeros([1, 1], dtype=cupy.float32)
+    dummy64 = cupy.zeros([1, 1], dtype=cupy.float64)
+
+    x_arr, y_arr = cupy.meshgrid(x_i, y_i)
+
+    xy_shape = x_arr.shape
+
+    x_inp = x_arr.flatten()
+    y_inp = y_arr.flatten()
+
+    line_len = cupy.zeros(x_inp.shape, cupy.float64)
+    N = cupy.array([x_inp.shape[0]], cupy.ulonglong)
+    s_len = cupy.array([1. / 16.], cupy.float32)
+    tol_coef = cupy.array([cupy.sqrt(0.1)], cupy.float32)
+
+    inp_norm = cupy.array([0, 0, 1.], cupy.float32)
+
+    twist_all = cupy.zeros(x_inp.shape, cupy.float64)
+
+    blck = (128, 1, 1)
+    grd = (28, 1)
+    cupy.cuda.stream.get_current_stream().synchronize()
+
+    Qube = np.zeros([xy_shape[1], xy_shape[0], z_i.shape[0]], dtype=np.float32)
+    Twube = np.zeros([xy_shape[1], xy_shape[0], z_i.shape[0]], dtype=np.float32)
+    Liube = np.zeros([xy_shape[1], xy_shape[0], z_i.shape[0]], dtype=np.float32)
+
+    pinned_mempool = cupy.get_default_pinned_memory_pool()
+    for idx_pos_z, z_pos in tqdm(enumerate(z_i), total=z_i.shape[0], desc='Compute Q and T'):
+
+        x_inp = x_arr.flatten()
+        y_inp = y_arr.flatten()
+        z_inp = (cupy.zeros_like(x_inp, cupy.float32) + z_pos)
+
+        (x_start, y_start, z_start, x_end, y_end, z_end,
+         Bx_start, By_start, Bz_start, Bx_end, By_end, Bz_end, Bx_inp, By_inp, Bz_inp
+         ) = [cupy.zeros(x_inp.shape, dtype=cupy.float32) for _ in range(15)]
+
+        (B_flag, flag_start, flag_end) = [cupy.zeros(x_inp.shape, dtype=cupy.int32) for _ in range(3)]
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+        # run the big calculation
+        trace_all_b_line(blck, grd,
+                         (Bx_gpu, By_gpu, Bz_gpu, BshapeN,
+                          Jx_gpu, Jy_gpu, Jz_gpu, twist_all, flag_twist,
+                          x_inp, y_inp, z_inp, inp_norm,
+                          x_start, y_start, z_start, flag_start,
+                          x_end, y_end, z_end, flag_end,
+                          Bx_inp, By_inp, Bz_inp, B_flag,
+                          Bx_start, By_start, Bz_start,
+                          Bx_end, By_end, Bz_end,
+                          s_len, N, line_len, tol_coef))
+        cupy.cuda.stream.get_current_stream().synchronize()
+        (x_end_arr, y_end_arr, z_end_arr, flag_end_arr,
+         x_start_arr, y_start_arr, z_start_arr, flag_start_arr,
+         Bx_in_arr, By_in_arr, Bz_in_arr,
+         Bx_out_arr, By_out_arr, Bz_out_arr,
+         Bx_0_arr, By_0_arr, Bz_0_arr, B_flag_arr, twist_all_arr,
+         line_len_arr) = FastQSL.ResReshape(xy_shape,
+                                            x_end, y_end, z_end, flag_end,
+                                            x_start, y_start, z_start, flag_start,
+                                            Bx_start, By_start, Bz_start,
+                                            Bx_end, By_end, Bz_end,
+                                            Bx_inp, By_inp, Bz_inp, B_flag, twist_all, line_len)
+
+        if z_pos < 1e-5:
+            B0cube = Bz_0_arr.get()
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+        Q = FastQSL.QCalcPlane(x_end_arr, y_end_arr, z_end_arr, flag_end_arr,
+                               x_start_arr, y_start_arr, z_start_arr, flag_start_arr,
+                               Bx_in_arr, By_in_arr, Bz_in_arr,
+                               Bx_out_arr, By_out_arr, Bz_out_arr,
+                               Bx_0_arr, By_0_arr, Bz_0_arr,
+                               B_flag_arr, stride_step)
+        stride_this = cupy.float32(1. / interp_ratio) / 8
+
+        (cut_inp_x, cut_inp_y, cut_inp_z,
+         cut_start_x, cut_start_y, cut_start_z, flag_cut_start,
+         cut_end_x, cut_end_y, cut_end_z, flag_cut_end,
+         Bx_inp_cut, By_inp_cut, Bz_inp_cut, B_flag_cut,
+         Bx_start_cut, By_start_cut, Bz_start_cut,
+         Bx_end_cut, By_end_cut, Bz_end_cut,
+         N_cut, line_len_cut, Bz0_start, Bz0_end) = FastQSL.CookPseudoLine(
+            x_end_arr, y_end_arr, z_end_arr, flag_end_arr,
+            x_start_arr, y_start_arr, z_start_arr, flag_start_arr,
+            Bx_in_arr, By_in_arr, Bz_in_arr,
+            Bx_out_arr, By_out_arr, Bz_out_arr,
+            Bz_0_arr, B_flag_arr, stride_this)
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+        trace_all_b_line(blck, grd,
+                         (Bx_gpu, By_gpu, Bz_gpu, BshapeN,
+                          dummy, dummy, dummy, dummy64, flag_twist_false,
+                          cut_inp_x, cut_inp_y, cut_inp_z, inp_norm,
+                          cut_start_x, cut_start_y, cut_start_z, flag_cut_start,
+                          cut_end_x, cut_end_y, cut_end_z, flag_cut_end,
+                          Bx_inp_cut, By_inp_cut, Bz_inp_cut, B_flag_cut,
+                          Bx_start_cut, By_start_cut, Bz_start_cut,
+                          Bx_end_cut, By_end_cut, Bz_end_cut,
+                          s_len, N_cut, line_len_cut, tol_coef * .1))
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+        (X1, Y1, X2, Y2) = [cupy.zeros(cut_inp_x.shape, dtype=cupy.float32)
+                            for _ in range(4)]
+
+        idx_Z1_cut = (flag_cut_start - 1) // 2 == 2
+        idx_Z2_cut = (flag_cut_end - 1) // 2 == 2
+        idx_Y1_cut = (flag_cut_start - 1) // 2 == 1
+        idx_Y2_cut = (flag_cut_end - 1) // 2 == 1
+        idx_X1_cut = (flag_cut_start - 1) // 2 == 0
+        idx_X2_cut = (flag_cut_end - 1) // 2 == 0
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+        # Z plane
+        X1[idx_Z1_cut] = cut_start_x[idx_Z1_cut]
+        Y1[idx_Z1_cut] = cut_start_y[idx_Z1_cut]
+        X2[idx_Z2_cut] = cut_end_x[idx_Z2_cut]
+        Y2[idx_Z2_cut] = cut_end_y[idx_Z2_cut]
+        # Y plane
+        X1[idx_Y1_cut] = cut_start_z[idx_Y1_cut]
+        Y1[idx_Y1_cut] = cut_start_x[idx_Y1_cut]
+        X2[idx_Y2_cut] = cut_end_z[idx_Y2_cut]
+        Y2[idx_Y2_cut] = cut_end_x[idx_Y2_cut]
+        # X plane
+        X1[idx_X1_cut] = cut_start_y[idx_X1_cut]
+        Y1[idx_X1_cut] = cut_start_z[idx_X1_cut]
+        X2[idx_X2_cut] = cut_end_y[idx_X2_cut]
+        Y2[idx_X2_cut] = cut_end_z[idx_X2_cut]
+
+        dx2xc = X2[0::4] - X2[2::4]
+        dx2yc = X2[1::4] - X2[3::4]
+        dy2xc = Y2[0::4] - Y2[2::4]
+        dy2yc = Y2[1::4] - Y2[3::4]
+        dx1xc = X1[0::4] - X1[2::4]
+        dx1yc = X1[1::4] - X1[3::4]
+        dy1xc = Y1[0::4] - Y1[2::4]
+        dy1yc = Y1[1::4] - Y1[3::4]
+        a_cut = (dx2xc * dy1yc - dx2yc * dy1xc)
+        b_cut = (dx2yc * dx1xc - dx2xc * dx1yc)
+        c_cut = (dy2xc * dy1yc - dy2yc * dy1xc)
+        d_cut = (dy2yc * dx1xc - dy2xc * dx1yc)
+
+        bnr_cut = cupy.abs(Bz0_end) / (cupy.abs(Bz0_start)) * ((1 / stride_this / 2) ** 4)
+        Qcut = (a_cut ** 2 + b_cut ** 2 + c_cut ** 2 + d_cut ** 2) * bnr_cut
+        Qcut[cupy.where(Qcut < 1.0)] = 1.0
+
+        Q_all = cupy.zeros(x_end_arr.shape, dtype=cupy.float32)
+        Q_all[1:-1, 1:-1] = Q
+        Q_all[B_flag_arr == 1] = Qcut
+
+        Qube[:, :, idx_pos_z] = Q_all.T.get()
+        Twube[:, :, idx_pos_z] = twist_all_arr.T.get()
+        Liube[:, :, idx_pos_z] = line_len_arr.T.get()
+
+        cupy.cuda.stream.get_current_stream().synchronize()
+
+        (x_start, y_start, z_start, x_end, y_end, z_end, z_inp,
+         Bx_start, By_start, Bz_start, Bx_end, By_end, Bz_end, Bx_inp, By_inp, Bz_inp
+         ) = [None for _ in range(16)]
+        (B_flag, flag_start, flag_end) = [None for _ in range(3)]
+        pinned_mempool.free_all_blocks()
+
+    Qube = Qube[::interp_ratio, ::interp_ratio, ::interp_ratio]
+    Twube = Twube[::interp_ratio, ::interp_ratio, ::interp_ratio]
+
+    base_name = os.path.basename(nf2_path).split('.')[0]
+    save_vtk(os.path.join(result_path, f"{base_name}.vtk"),
+             scalars={"Q": Qube, "Twist": Twube},
+             vectors={"B": output["b"]})
diff --git a/nf2/evaluation/series.py b/nf2/evaluation/series.py
deleted file mode 100644
index edd54ea..0000000
--- a/nf2/evaluation/series.py
+++ /dev/null
@@ -1,173 +0,0 @@
-import argparse
-import glob
-import os
-
-import imageio
-import numpy as np
-import torch.cuda
-from dateutil.parser import parse
-from matplotlib import pyplot as plt, cm
-from matplotlib.colors import Normalize
-from matplotlib.dates import date2num, DateFormatter
-from sunpy.net import Fido
-from sunpy.net import attrs as a
-from tqdm import tqdm
-
-from nf2.evaluation.energy import get_free_mag_energy
-from nf2.evaluation.metric import energy, curl, vector_norm, divergence, weighted_theta
-from nf2.evaluation.unpack import load_cube
-
-
-def evaluate_nf2(nf2_file, z, cm_per_pixel, strides, batch_size):
-    b = load_cube(nf2_file, progress=False, z=z, strides=strides, batch_size=batch_size)
-    j = curl(b)
-    jxb = np.cross(j, b, axis=-1)
-    me = energy(b)
-    free_me = get_free_mag_energy(b, progress=False)
-    theta = weighted_theta(b, j)
-
-    result = {
-        'date': parse(os.path.basename(nf2_file).split('.')[0][:-4].replace('_', 'T')),
-        # torch.load(nf2_file)['meta_data']['DATE-OBS'],
-        # bottom boundary
-        'b_0': b[:, :, 0, 2],
-        # integrated quantities
-        'total_energy': me.sum() * cm_per_pixel ** 3,
-        'total_free_energy': free_me.sum() * cm_per_pixel ** 3,
-        'total_div': (np.abs(divergence(b)) / vector_norm(b)).mean(),
-        'total_jxb': vector_norm(jxb).mean(),
-        'theta': theta,
-        # heigth distribution
-        'height_free_energy': free_me.mean((0, 1)),
-        # maps
-        'j_map': vector_norm(j).sum(2),
-        'energy_map': me.sum(2) * cm_per_pixel,
-        'free_energy_map': free_me.sum(2) * cm_per_pixel,
-        'jxb_map': vector_norm(jxb).sum(2),
-    }
-    return result
-
-
-def evaluate_nf2_series(nf2_paths, z, cm_per_pixel, strides, batch_size):
-    results = []
-    for nf2_file in tqdm(nf2_paths):
-        results.append(evaluate_nf2(nf2_file, z, cm_per_pixel, strides, batch_size))
-    # concatenate dicts to dict of lists
-    results = {k: [r[k] for r in results] for k in results[0].keys()}
-    return results
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
-    parser.add_argument('nf2_path', type=str, help='path to the directory of the NF2 files')
-    parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
-    parser.add_argument('--strides', type=int, help='downsampling of the volume', required=False, default=1)
-    parser.add_argument('--add_flares', action='store_true', help='query for flares and plot', required=False,
-                        default=None)
-    args = parser.parse_args()
-
-    nf2_files = sorted(glob.glob(os.path.join(args.nf2_path, '*.nf2')))
-    result_path = args.result_path
-    os.makedirs(result_path, exist_ok=True)
-
-    batch_size = int(1e5 * torch.cuda.device_count())
-
-    # load simulation scaling
-    Mm_per_pix = .72  # torch.load(nf2_files[0])['Mm_per_pix']
-    z_pixels = int(np.ceil(20 / (Mm_per_pix)))  # 20 Mm --> pixels
-
-    # adjust scale to strides
-    Mm_per_pix *= args.strides
-    cm_per_pix = (Mm_per_pix * 1e8)
-
-    # evaluate series
-    series_results = evaluate_nf2_series(nf2_files, z_pixels, cm_per_pix, args.strides, batch_size)
-
-    # plot integrated quantities
-    x_dates = date2num(series_results['date'])
-    date_format = DateFormatter('%d-%H:%M')
-
-    fig, full_axs = plt.subplots(5, 2, figsize=(8, 12), gridspec_kw={"width_ratios": [1, 0.05]})
-    axs = full_axs[:, 0]
-    [ax.axis('off') for ax in full_axs[:, 1]]
-    # make date axis
-    for ax in axs:
-        ax.xaxis_date()
-        ax.set_xlim(x_dates[0], x_dates[-1])
-        ax.xaxis.set_major_formatter(date_format)
-
-    fig.autofmt_xdate()
-
-    ax = axs[0]
-    ax.plot(x_dates, np.array(series_results['total_energy']) * 1e-32)
-    ax.set_ylabel('Energy\n[$10^{32}$ erg]')
-
-    ax = axs[1]
-    ax.plot(x_dates, np.array(series_results['total_free_energy']) * 1e-32)
-    ax.set_ylabel('Free Energy\n[$10^{32}$ erg]')
-
-    ax = axs[2]
-    dt = np.diff(x_dates)[0] / 2
-    dz = Mm_per_pix / 2
-    free_energy_distribution = np.array(series_results['height_free_energy']).T
-    max_height = free_energy_distribution.shape[0] * Mm_per_pix
-    mpb = ax.imshow(free_energy_distribution,  # average
-                    extent=(x_dates[0] - dt, x_dates[-1] + dt, -dz, max_height + dz), aspect='auto', origin='lower',
-                    cmap=cm.get_cmap('jet'), vmin=0)
-    ax.set_ylabel('Altitude\n[Mm]')
-    ax.set_ylim(0, 20)
-    # add colorbar
-    full_axs[2, 1].axis('on')
-    cbar = fig.colorbar(mpb, cax=full_axs[2, 1], label='Free Energy Density\n' + r'[erg $cm^{-3}]$')
-    cbar.formatter.set_powerlimits((0, 0))
-    cbar.set_ticks([2e2, 4e2, 6e2, 8e2])
-    full_axs[2, 1].remove()
-
-    ax = axs[3]
-    ax.plot(x_dates, np.array(series_results['total_div']))
-    ax.set_ylabel(r'$|\mathbf{\nabla} \cdot B| / |B|$' + '\n[1/pixel]')
-
-    ax = axs[4]
-    ax.plot(x_dates, np.array(series_results['theta']))
-    ax.set_ylabel(r'$\theta_J$' + '\n[deg]')
-
-    if args.add_flares:
-        flares = Fido.search(a.Time(min(series_results['date']), max(series_results['date'])),
-                             a.hek.EventType("FL"),
-                             a.hek.OBS.Observatory == "GOES")["hek"]
-        goes_mapping = {c: 10 ** (i) for i, c in enumerate(['B', 'C', 'M', 'X'])}
-        for t, goes_class in zip(flares['event_peaktime'], flares['fl_goescls']):
-            flare_intensity = np.log10(float(goes_class[1:]) * goes_mapping[goes_class[0]])
-            if flare_intensity >= np.log10(1 * 10 ** 2):
-                [ax.axvline(x=date2num(t.datetime), linestyle='dotted', c='black') for ax in axs]
-            if flare_intensity >= np.log10(1 * 10 ** 3):
-                [ax.axvline(x=date2num(t.datetime), linestyle='dotted', c='red') for ax in axs]
-
-    plt.tight_layout()
-    plt.savefig(os.path.join(result_path, 'series.jpg'))
-    plt.close()
-
-    # create video of maps
-    # j_map
-    j_maps = np.array(series_results['j_map'])
-    images = [cm.get_cmap('viridis')(Normalize(vmin=0, vmax=j_maps.max())(j_map.T)) for j_map in j_maps]
-    images = np.flip(images, axis=1)
-    imageio.mimsave(os.path.join(result_path, 'j_maps.mp4'), images)
-    # free_energy_map
-    free_energy_maps = np.array(series_results['free_energy_map'])
-    images = [cm.get_cmap('jet')(Normalize(vmin=0, vmax=free_energy_maps.max())(free_energy_map.T)) for free_energy_map
-              in free_energy_maps]
-    images = np.flip(images, axis=1)
-    imageio.mimsave(os.path.join(result_path, 'free_energy_maps.mp4'), images)
-    # free_energy_change_map
-    free_energy_change_maps = np.gradient(np.array(series_results['free_energy_map']), axis=0)
-    v_min_max = np.max(np.abs(free_energy_change_maps))
-    images = [cm.get_cmap('seismic')(Normalize(vmin=-v_min_max, vmax=v_min_max)(free_energy_change_map.T)) for
-              free_energy_change_map in free_energy_change_maps]
-    images = np.flip(images, axis=1)
-    imageio.mimsave(os.path.join(result_path, 'free_energy_change_maps.mp4'), images)
-    # jxb_map
-    jxb_maps = np.array(series_results['jxb_map'])
-    images = [cm.get_cmap('plasma')(Normalize(vmin=0, vmax=jxb_maps.max())(jxb_map.T)) for jxb_map in jxb_maps]
-    images = np.flip(images, axis=1)
-    imageio.mimsave(os.path.join(result_path, 'jxb_maps.mp4'), images)
diff --git a/nf2/evaluation/sharp/__init__.py b/nf2/evaluation/sharp/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/evaluation/sharp/convert_series.py b/nf2/evaluation/sharp/convert_series.py
new file mode 100644
index 0000000..c0d39d3
--- /dev/null
+++ b/nf2/evaluation/sharp/convert_series.py
@@ -0,0 +1,230 @@
+import argparse
+import glob
+import os
+import pickle
+from copy import copy
+
+import numpy as np
+from astropy import units as u
+from astropy.io import fits
+from dateutil.parser import parse
+from matplotlib import pyplot as plt, cm
+from matplotlib.colors import Normalize, LogNorm
+from matplotlib.dates import date2num, DateFormatter
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from sunpy.net import Fido
+from sunpy.net import attrs as a
+from tqdm import tqdm
+
+from nf2.evaluation.energy import get_free_mag_energy
+from nf2.evaluation.metric import energy, vector_norm, divergence, theta_J
+from nf2.evaluation.output import CartesianOutput
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--nf2_path', type=str, help='path to the directory of the NF2 files', required=True)
+    parser.add_argument('--euv_path', type=str, help='path to the directory of the EUV files', required=True)
+    parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+    parser.add_argument('--max_height', type=float, help='max height of the volume in Mm', required=False, default=20)
+    parser.add_argument('--Mm_per_pixel', type=float, help='Mm per pixel', required=False, default=0.72)
+    parser.add_argument('--add_flares', action='store_true', help='query for flares and plot', required=False,
+                        default=None)
+    args = parser.parse_args()
+
+    nf2_files = sorted(glob.glob(args.nf2_path))
+    result_path = args.result_path if args.result_path is not None else args.nf2_path
+    os.makedirs(result_path, exist_ok=True)
+
+    # evaluate series
+    series_results = convert_nf2_series(nf2_files, result_path,
+                                        height_range=[0, args.max_height],
+                                        Mm_per_pixel=args.Mm_per_pixel, batch_size=int(2 ** 14))
+
+    times, integrated_quantities, metrics, maps, height_distribution, wcs = load_results(series_results)
+
+    # plot integrated quantities
+    _plot_integrated_qunatities(times, integrated_quantities, height_distribution,
+                                result_path,
+                                args.Mm_per_pixel, add_flares=args.add_flares)
+
+    _plot_map_videos(times, maps, wcs, args.euv_path, result_path, args.Mm_per_pixel)
+
+
+def evaluate_nf2(nf2_file, **kwargs):
+    out = CartesianOutput(nf2_file)
+    res = out.load_cube(**kwargs)
+
+    b = res['b']
+    j = res['j']
+    a = res['a']
+    Mm_per_pixel = res['Mm_per_pixel'] * u.Mm
+
+    me = energy(b).value * (u.erg * u.cm ** -3)  # erg = G^2 cm^3
+    free_me = get_free_mag_energy(b.to_value(u.G), progress=False) * (u.erg * u.cm ** -3)
+    theta = theta_J(b.value, j.value)
+    magnetic_helicity = (a * b).sum(-1)
+    current_helicity = (b * j).sum(-1)
+    jxb = np.cross(j, b, axis=-1)
+
+    result = {
+        'time': out.time,
+        'integrated_quantities': {
+            'energy': me.sum() * Mm_per_pixel ** 3,
+            'free_energy': free_me.sum() * Mm_per_pixel ** 3,
+            'current_helicity': current_helicity.sum() * Mm_per_pixel ** 3,
+            'magnetic_helicity': magnetic_helicity.sum() * Mm_per_pixel ** 3,
+        },
+        'metrics': {
+            'divergence': (np.abs(divergence(b)) / vector_norm(b)).mean(),
+            'jxb': vector_norm(jxb).mean(),
+            'theta': theta
+        },
+        'maps': {
+            'b_0': b[:, :, 0, 2],  # bottom boundary
+            'current_density_map': vector_norm(j).sum(2) * Mm_per_pixel,
+            'energy_map': me.sum(2) * Mm_per_pixel,
+            'free_energy_map': free_me.sum(2) * Mm_per_pixel,
+            'jxb_map': vector_norm(jxb).sum(2) * Mm_per_pixel,
+        },
+        'height_distribution': {
+            'height_free_energy': free_me.mean((0, 1)) * Mm_per_pixel ** 2,
+            'height_current_density': vector_norm(j).mean((0, 1)) * Mm_per_pixel ** 2,
+        },
+        'info': {
+            'data_config': out.data_config,
+            'Mm_per_pixel': res['Mm_per_pixel'],
+            'wcs': out.wcs
+        }
+    }
+    return result
+
+
+def convert_nf2_series(nf2_paths, result_path, **kwargs):
+    series_result_path = os.path.join(result_path, 'pkl')
+    os.makedirs(series_result_path, exist_ok=True)
+    results = []
+    for nf2_file in tqdm(nf2_paths):
+        save_file = os.path.join(series_result_path, os.path.basename(nf2_file).replace('.nf2', '.pkl'))
+        if os.path.exists(save_file):
+            results.append(save_file)
+            continue
+        result = evaluate_nf2(nf2_file, **kwargs)
+        with open(save_file, 'wb') as f:
+            pickle.dump(result, f)
+        results.append(save_file)
+
+    return results
+
+
+def _plot_map_videos(times, maps, wcs, euv_path, result_path, Mm_per_pixel):
+    euv_files = np.array(sorted(glob.glob(euv_path)))
+    euv_dates = np.array([parse(fits.getheader(f, 1)['DATE-OBS']) for f in euv_files])
+
+    video_path = os.path.join(result_path, 'video')
+    os.makedirs(video_path, exist_ok=True)
+
+    j_norm = LogNorm()
+    free_energy_norm = LogNorm(1e9, 1e14)
+
+    for i, time in enumerate(times):
+        print(time)
+        euv_file = euv_files[np.argmin(np.abs(euv_dates - time))]
+        euv_map = Map(euv_file)
+        exposure = euv_map.exposure_time.to_value(u.s)
+        euv_map = euv_map.reproject_to(wcs[i][0])
+
+        b_0 = maps['b_0'][i].value
+        current_density_map = maps['current_density_map'][i].to_value(u.G * u.cm * u.s ** -1)
+        free_energy_map = maps['free_energy_map'][i].to_value(u.erg * u.cm ** -2)
+
+        fig, axs = plt.subplots(2, 2, figsize=(10, 6))
+
+        ax = axs[0, 0]
+
+        extent = np.array([0, b_0.shape[0], 0, b_0.shape[1]]) * Mm_per_pixel
+        im = ax.imshow(b_0.T, origin='lower', cmap='gray', extent=extent, norm=Normalize(vmin=-2000, vmax=2000))
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="3%", pad=0.05)
+        cbar = fig.colorbar(im, cax=cax, label='[G]')
+        ax.set_title('Magnetic Field - B$_z$')
+
+        ax = axs[0, 1]
+        cm = copy(plt.get_cmap('sdoaia131'))
+        cm.set_bad('black')
+        im = ax.imshow(euv_map.data / exposure, origin='lower', cmap=cm, extent=extent, norm=LogNorm(vmin=5, vmax=1e3))
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="3%", pad=0.05)
+        cbar = fig.colorbar(im, cax=cax, label='[DN / s]')
+        ax.set_title('EUV - AIA 131')
+
+        ax = axs[1, 0]
+
+        im = ax.imshow(current_density_map.T, origin='lower', cmap='inferno', extent=extent, norm=j_norm)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="3%", pad=0.05)
+        cbar = fig.colorbar(im, cax=cax, label='[G cm / s]')
+        ax.set_title('Current Density - |J|')
+
+        ax = axs[1, 1]
+        im = ax.imshow(free_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=free_energy_norm)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="3%", pad=0.05)
+        cbar = fig.colorbar(im, cax=cax, label='[erg / cm$^2$]')
+        ax.set_title('Free Magnetic Energy')
+
+        [ax.set_xlabel('X [Mm]') for ax in axs[1, :]]
+        [ax.set_ylabel('Y [Mm]') for ax in axs[:, 0]]
+
+        fig.suptitle(f'Time: {time.isoformat(" ", timespec="minutes")}')
+
+        plt.tight_layout()
+        plt.savefig(os.path.join(video_path, f'{time.isoformat("T", timespec="minutes")}.jpg'), dpi=300)
+        plt.close()
+
+
+def load_results(series_results):
+    series_results = series_results if isinstance(series_results, list) else sorted(glob.glob(series_results))
+    # load files
+    integrated_quantities = {'energy': [], 'free_energy': [], 'current_helicity': [], 'magnetic_helicity': []}
+    metrics = {'divergence': [], 'jxb': [], 'theta': []}
+    maps = {'b_0': [], 'current_density_map': [], 'energy_map': [], 'free_energy_map': [], 'jxb_map': []}
+    height_distribution = {'height_free_energy': [], 'height_current_density': []}
+    times = []
+    wcs = []
+    Mm_per_pixel = None
+    for f in series_results:
+        with open(f, 'rb') as file:
+            data = pickle.load(file)
+            integrated_quantities['energy'].append(data['integrated_quantities']['energy'])
+            integrated_quantities['free_energy'].append(data['integrated_quantities']['free_energy'])
+            integrated_quantities['current_helicity'].append(data['integrated_quantities']['current_helicity'])
+            integrated_quantities['magnetic_helicity'].append(data['integrated_quantities']['magnetic_helicity'])
+            metrics['divergence'].append(data['metrics']['divergence'])
+            metrics['jxb'].append(data['metrics']['jxb'])
+            metrics['theta'].append(data['metrics']['theta'])
+            maps['b_0'].append(data['maps']['b_0'])
+            maps['current_density_map'].append(data['maps']['current_density_map'])
+            maps['energy_map'].append(data['maps']['energy_map'])
+            maps['free_energy_map'].append(data['maps']['free_energy_map'])
+            maps['jxb_map'].append(data['maps']['jxb_map'])
+            height_distribution['height_free_energy'].append(data['height_distribution']['height_free_energy'])
+            height_distribution['height_current_density'].append(data['height_distribution']['height_current_density'])
+            times.append(data['time'])
+            wcs.append(data['info']['wcs'])
+            Mm_per_pixel = data['info']['Mm_per_pixel']
+
+    integrated_quantities = {k: np.stack(v) for k, v in integrated_quantities.items()}
+    metrics = {k: np.stack(v) for k, v in metrics.items()}
+    maps = {k: np.stack(v) for k, v in maps.items()}
+    height_distribution = {k: np.stack(v) for k, v in height_distribution.items()}
+    times = np.array(times)
+    wcs = np.array(wcs)
+
+    return {'times': times, 'integrated_quantities': integrated_quantities, 'metrics': metrics, 'maps': maps,
+            'height_distribution': height_distribution, 'wcs': wcs, 'Mm_per_pixel': Mm_per_pixel}
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/evaluation/sharp/ensemble.py b/nf2/evaluation/sharp/ensemble.py
new file mode 100644
index 0000000..4e2da59
--- /dev/null
+++ b/nf2/evaluation/sharp/ensemble.py
@@ -0,0 +1,128 @@
+import argparse
+import glob
+import os.path
+
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.output import CartesianOutput
+from astropy import units as u
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--nf2_path', type=str, help='path to the directory of the NF2 files. Uses glob search pattern. (<<your path>>/**/extrapolation_result.nf2)', required=True)
+parser.add_argument('--ref_path', type=str, help='path to the reference map', required=False, default=None)
+parser.add_argument('--ref_euv_path', type=str, help='path to the reference map', required=False, default=None)
+parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+parser.add_argument('--Mm_per_pixel', type=float, help='Mm per pixel', required=False, default=0.72)
+args = parser.parse_args()
+
+runs = sorted(glob.glob(args.nf2_path, recursive=True))
+ref_map = Map(args.ref_path)
+
+ref_euv_map = Map('/glade/work/rjarolim/data/nf2/13664/euv/aia.lev1_euv_12s.2024-05-09T000007Z.304.image_lev1.fits')
+euv_cmap = ref_euv_map.plot_settings['cmap']
+ref_euv_map = Map(ref_euv_map.data / ref_euv_map.exposure_time.to(u.s), ref_euv_map.meta)
+ref_euv_map = ref_euv_map.reproject_to(ref_map.wcs)
+
+b_ensemble = []
+j_ensemble = []
+for f in tqdm(runs, desc='Loading Ensemble'):
+    model = CartesianOutput(f)
+    out = model.load_cube(Mm_per_pixel=args.Mm_per_pixel)
+    b_ensemble.append(out['b'])
+    j_ensemble.append(out['j'])
+
+Mm_per_pixel = out['Mm_per_pixel']
+
+b_ensemble = np.stack(b_ensemble).to_value(u.G)
+j_ensemble = np.stack(j_ensemble).to_value(u.G / u.s)
+
+mean_b_vector = np.mean(b_ensemble, axis=0, keepdims=True)
+b_uncertainty = np.mean(np.linalg.norm(b_ensemble - mean_b_vector, axis=-1) ** 2, axis=0) ** 0.5
+
+mean_j_vector = np.mean(j_ensemble, axis=0, keepdims=True)
+j_uncertainty = np.mean(np.linalg.norm(j_ensemble - mean_j_vector, axis=-1) ** 2, axis=0) ** 0.5
+
+extent_xy = [0, b_uncertainty.shape[0] * Mm_per_pixel, 0, b_uncertainty.shape[1] * Mm_per_pixel]
+extent_xz = [0, b_uncertainty.shape[0] * Mm_per_pixel, 0, b_uncertainty.shape[2] * Mm_per_pixel]
+
+fig, ax = plt.subplots(1, 1, figsize=(4, 3))
+
+im = ax.imshow(ref_map.data, cmap='gray', origin='lower', vmin=-2000, vmax=2000, extent=extent_xy)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='3%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='[G]')
+ax.set_title('Magnetic Field - B$_z$')
+ax.set_xlabel('X [Mm]')
+ax.set_ylabel('Y [Mm]')
+
+plt.tight_layout()
+plt.savefig('/glade/work/rjarolim/nf2/sharp/13664/ensemble/mag_reference.png', dpi=300, transparent=True)
+plt.close()
+
+# ax = axs[1]
+# im = ax.imshow(ref_euv_map.data, cmap=euv_cmap, origin='lower', extent=extent_xy, norm=LogNorm(vmin=1))
+# divider = make_axes_locatable(ax)
+# cax = divider.append_axes('right', size='3%', pad=0.05)
+# fig.colorbar(im, cax=cax, orientation='vertical', label='[DN/s]')
+# ax.set_title('SDO/AIA 131 $\AA$')
+# ax.set_xlabel('X [Mm]')
+# ax.set_ylabel('Y [Mm]')
+
+fig, axs = plt.subplots(2, 1, figsize=(4, 4))
+
+ax = axs[0]
+im = ax.imshow(b_uncertainty.mean(2).T, cmap='viridis', origin='lower', extent=extent_xy)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='3%', pad=0.05)
+ax.contour(ref_map.data, levels=[-1500, 1500], colors=['black', 'white'], alpha=0.6, extent=extent_xy)
+fig.colorbar(im, cax=cax, orientation='vertical', label='[G]')
+ax.set_title('Uncertainty')
+ax.set_xlabel('X [Mm]')
+ax.set_ylabel('Y [Mm]')
+
+ax = axs[1]
+im = ax.imshow(b_uncertainty.mean(1).T, cmap='viridis', origin='lower', extent=extent_xz)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='3%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='[G]')
+ax.set_xlabel('X [Mm]')
+ax.set_ylabel('Z [Mm]')
+
+
+plt.tight_layout()
+plt.savefig('/glade/work/rjarolim/nf2/sharp/13664/ensemble/overview.png', dpi=300, transparent=True)
+plt.close()
+
+j_norm = LogNorm()
+
+fig, axs = plt.subplots(2, 3, figsize=(8, 3))
+
+[ax.set_ylabel('Y [Mm]') for ax in axs[:, 0]]
+[ax.set_xlabel('X [Mm]') for ax in axs[1, :]]
+
+axs = np.ravel(axs)
+for i, (f, ax) in enumerate(zip(runs, axs[:-1])):
+    current_density = np.linalg.norm(j_ensemble[0], axis=-1).sum(2) * Mm_per_pixel
+    im = ax.imshow(current_density.T, cmap='inferno', origin='lower', extent=extent_xy, norm=j_norm)
+    ax.set_title(f'Run {i + 1}')
+
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='3%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='[G cm / s]')
+
+ax = axs[-1]
+current_density_std = np.std(np.linalg.norm(j_ensemble, axis=-1).sum(3) * Mm_per_pixel, axis=0)
+im = ax.imshow(current_density_std.T, cmap='viridis', origin='lower', extent=extent_xy, norm=j_norm)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='3%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='[G cm / s]')
+ax.set_title('Ensemble Std')
+
+plt.tight_layout()
+plt.savefig(os.path.join(args.result_path, 'j_overview.png'), dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sharp/flares.py b/nf2/evaluation/sharp/flares.py
new file mode 100644
index 0000000..d2c654c
--- /dev/null
+++ b/nf2/evaluation/sharp/flares.py
@@ -0,0 +1,191 @@
+import argparse
+import glob
+import os
+from datetime import timedelta
+from multiprocessing import Pool
+
+import drms
+import matplotlib.pyplot as plt
+import numpy as np
+from astropy import units as u
+from dateutil.parser import parse
+from matplotlib.colors import LogNorm
+from matplotlib.dates import DateFormatter
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+
+from nf2.data.download import download_euv
+from nf2.evaluation.energy import get_free_mag_energy
+from nf2.evaluation.metric import energy
+from nf2.evaluation.output import CartesianOutput
+
+
+class _F:
+    def __init__(self, ref_wcs):
+        self.ref_wcs = ref_wcs
+
+    def func(self, file):
+        s_map = Map(file)
+        exposure = s_map.exposure_time.to_value(u.s)
+        reprojected_map = s_map.reproject_to(self.ref_wcs)
+        return Map(reprojected_map.data / exposure, self.ref_wcs)
+
+
+def get_integrated_euv_map(euv_files, ref_wcs):
+    with Pool(os.cpu_count()) as p:
+        reprojected_maps = p.map(_F(ref_wcs).func, euv_files)
+    integrated_euv = np.array([m.data for m in reprojected_maps]).sum(0)
+    time = (Map(euv_files[-1]).date.datetime - Map(euv_files[0]).date.datetime).total_seconds()
+    euv_map = Map(integrated_euv * time, ref_wcs)
+    return euv_map
+
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--nf2_path', type=str, help='path to the directory of the NF2 files')
+parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+parser.add_argument('--email', type=str, help='email for the DRMS client', required=True)
+parser.add_argument('--date_range', nargs=3, type=str, required=True)
+args = parser.parse_args()
+
+nf2_files = sorted(glob.glob(args.nf2_path))
+result_path = args.result_path
+os.makedirs(result_path, exist_ok=True)
+
+dates = [parse(os.path.basename(nf2_file).split('.')[0][:-4].replace('_', 'T')) for nf2_file in nf2_files]
+dates = np.array(dates)
+
+start_time = parse(args.date_range[0])
+end_time = parse(args.date_range[1])
+peak_time = parse(args.date_range[2])
+
+client = drms.Client(email=args.email, verbose=True)
+euv_files = download_euv(start_time=start_time, end_time=end_time,
+                         dir=result_path, client=client, channel=94)
+
+filter_dates = (dates > (start_time - timedelta(minutes=12))) & \
+               (dates < (end_time + timedelta(minutes=12)))
+# filter_dates = (dates > start_time) & (dates < end_time)
+flare_nf2_files = np.array(nf2_files)[filter_dates]
+
+model_1 = CartesianOutput(flare_nf2_files[0])
+out_1 = model_1.load_cube(height_range=[0, 20], Mm_per_pixel=0.72)
+free_energy_1 = get_free_mag_energy(out_1['b'].value, progress=False) * (u.erg / u.cm ** 3)
+energy_1 = energy(out_1['b'].value) * (u.erg / u.cm ** 3)
+
+model_2 = CartesianOutput(flare_nf2_files[-1])
+out_2 = model_2.load_cube(height_range=[0, 20], Mm_per_pixel=0.72)
+free_energy_2 = get_free_mag_energy(out_2['b'].value, progress=False) * (u.erg / u.cm ** 3)
+energy_2 = energy(out_2['b'].value) * (u.erg / u.cm ** 3)
+
+Mm_per_pix = out_1['Mm_per_pixel'] * u.Mm
+
+released_energy = -np.clip(energy_1 - energy_2, a_min=None, a_max=0)
+released_energy_map = (released_energy.sum(2) * Mm_per_pix).to_value(u.erg / u.cm ** 2)
+released_energy_map[released_energy_map < 1e11] = np.nan
+total_released_energy = (released_energy.sum() * Mm_per_pix ** 3).to_value(u.erg)
+
+energy_diff = (energy_2 - energy_1).sum() * Mm_per_pix ** 3
+energy_diff = energy_diff.to_value(u.erg)
+free_energy_diff = (free_energy_2 - free_energy_1).sum() * Mm_per_pix ** 3
+free_energy_diff = free_energy_diff.to_value(u.erg)
+
+with open(os.path.join(args.result_path, f'{start_time.isoformat("T", timespec="minutes")}.txt'), 'w') as f:
+    f.write(f'Total released energy: {total_released_energy:.2e} erg\n')
+    f.write(f'Total energy difference: {energy_diff:.2e} erg\n')
+    f.write(f'Total free energy difference: {free_energy_diff:.2e} erg\n')
+
+integrated_currents_1 = (np.linalg.norm(out_1['j'], axis=-1).sum(2) * Mm_per_pix).to_value(u.G * u.cm / u.s)
+integrated_currents_2 = (np.linalg.norm(out_2['j'], axis=-1).sum(2) * Mm_per_pix).to_value(u.G * u.cm / u.s)
+
+euv_map = get_integrated_euv_map(list(euv_files), model_1.wcs[0])
+[os.remove(f) for f in euv_files]
+
+extent = (0, released_energy_map.shape[0] * Mm_per_pix.to_value(u.Mm), 0,
+          released_energy_map.shape[1] * Mm_per_pix.to_value(u.Mm))
+
+fig, axs = plt.subplots(1, 3, figsize=(12, 3))
+
+ax = axs[0]
+im = ax.imshow(euv_map.data, origin='lower', cmap='sdoaia94', norm=LogNorm(np.percentile(euv_map.data, 10)),
+               extent=extent)
+# ax.set_title('Integrated EUV SDO/AIA 94 $\AA$')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="3%", pad=0.05)
+fig.colorbar(im, cax=cax, label='AIA 94 $\AA$ [DN]')
+ax.contour(released_energy_map.T, levels=[1e12], colors='red', linewidths=1, extent=extent)
+
+ax = axs[1]
+free_energy_map = (free_energy_2.sum(2) * Mm_per_pix).to_value(u.erg / u.cm ** 2)
+im = ax.imshow(free_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=LogNorm(1e11, 3e13))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="3%", pad=0.05)
+fig.colorbar(im, cax=cax, label='E$_{free}$ [erg / cm$^2$]')
+
+ax = axs[2]
+im = ax.imshow(released_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=LogNorm(1e11, 3e13))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="3%", pad=0.05)
+fig.colorbar(im, cax=cax, label='$\Delta$ E [erg / cm$^2$]')
+
+plt.tight_layout()
+plt.savefig(os.path.join(args.result_path, f'fl_{start_time.isoformat("T", timespec="minutes")}'), dpi=300,
+            transparent=True)
+plt.close()
+
+Mm_per_ds = model_1.Mm_per_ds
+cond = (dates > (start_time - timedelta(hours=2))) & \
+       (dates < (end_time + timedelta(minutes=60)))
+# filter_dates = (dates > start_time) & (dates < end_time)
+series_nf2_files = np.array(nf2_files)[cond]
+series_dates = dates[cond]
+
+energies = []
+free_energies = []
+for f in series_nf2_files:
+    model = CartesianOutput(f)
+    out = model.load_cube(height_range=[0, 20], Mm_per_pixel=0.72)
+    e_free = get_free_mag_energy(out['b'].to_value(u.G)) * (u.erg / u.cm ** 3)
+    e = energy(out['b'].to_value(u.G)) * (u.erg / u.cm ** 3)
+    #
+    total_e = e.sum() * (out['Mm_per_pixel'] * u.Mm) ** 3
+    total_e_free = e_free.sum() * (out['Mm_per_pixel'] * u.Mm) ** 3
+    #
+    energies.append(total_e.to_value(u.erg))
+    free_energies.append(total_e_free.to_value(u.erg))
+
+date_format = DateFormatter('%d-%H:%M')
+fig, ax1 = plt.subplots(1, 1, figsize=(7, 3))
+
+color = 'tab:blue'
+ax1.set_ylabel('E [10$^{32}$ erg]', fontdict={'size': 16}, color=color)
+l1 = ax1.plot(series_dates, np.array(energies) * 1e-32, color=color, label='Magnetic Energy')
+ax1.tick_params(axis='y', labelcolor=color)
+
+ax2 = ax1.twinx()  # instantiate a second Axes that shares the same x-axis
+
+color = 'tab:red'
+ax2.set_ylabel('E$_{free}$ [10$^{32}$ erg]', fontdict={'size': 16}, color=color)
+l2 = ax2.plot(series_dates, np.array(free_energies) * 1e-32, color=color, label='Free Magnetic Energy')
+ax2.tick_params(axis='y', labelcolor=color)
+
+ax1.set_xlim(series_dates[0], series_dates[-1])
+ax1.axvspan(xmin=start_time, xmax=end_time, color='orange', alpha=0.5)
+ax1.set_xticks([series_dates[1], peak_time, series_dates[-2]])
+ax1.xaxis.set_major_formatter(date_format)
+ax1.axvline(x=peak_time, color='black', linestyle='--')
+
+# legend in forground fancy and transparent
+ax1.legend(l1 + l2, [l.get_label() for l in l1 + l2],
+           loc='upper left', fancybox=True, framealpha=0.8)
+ax1.set_zorder(1)
+
+# change font size of ticks
+ax1.tick_params(axis='x', labelsize=16)
+ax1.tick_params(axis='y', labelsize=16)
+ax2.tick_params(axis='y', labelsize=16)
+
+fig.tight_layout()
+plt.savefig(os.path.join(args.result_path, f'energy_{start_time.isoformat("T", timespec="minutes")}'), dpi=300,
+            transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sharp/plot_integrated_quantities.py b/nf2/evaluation/sharp/plot_integrated_quantities.py
new file mode 100644
index 0000000..bdcb7ae
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_integrated_quantities.py
@@ -0,0 +1,140 @@
+import argparse
+import os
+from datetime import datetime
+
+import numpy as np
+import pandas as pd
+from matplotlib import pyplot as plt, cm
+from matplotlib.colors import LogNorm
+from matplotlib.dates import DateFormatter
+
+from astropy import units as u
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.net import Fido
+from sunpy.net import attrs as a
+
+from nf2.evaluation.sharp.convert_series import load_results
+
+from sunpy import timeseries as ts
+
+def main():
+    pass
+
+
+def _plot_integrated_qunatities(times, integrated_quantities, height_distribution, result_path, Mm_per_pixel):
+    pass
+
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--pkl_path', type=str, help='path to the directory with the converted pkl files.')
+parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+args = parser.parse_args()
+
+output = load_results(args.pkl_path)
+
+# plot integrated quantities
+# _plot_integrated_qunatities(output['times'], output['integrated_quantities'], output['height_distribution'],
+#                             args.result_path, output['Mm_per_pixel'])
+times = output['times']
+integrated_quantities = output['integrated_quantities']
+height_distribution = output['height_distribution']
+result_path = args.result_path
+Mm_per_pixel = output['Mm_per_pixel']
+
+free_energy = integrated_quantities['free_energy'].to_value(u.erg) * 1e-32
+#
+free_energy_distribution = height_distribution['height_free_energy'].to_value(u.erg * u.cm ** -1)
+df = pd.DataFrame(index=times, data={'idx': np.arange(len(times), dtype=int)})
+df = df.groupby(pd.Grouper(freq='12Min')).first()
+
+filled_energy_distribution = np.ones((len(df), *free_energy_distribution.shape[1:])) * np.nan
+cond = ~np.isnan(df['idx'].values)
+filled_energy_distribution[cond, :] = free_energy_distribution
+filled_times = df.index.to_pydatetime()
+
+result_goes = Fido.search(a.Time(times[0], times[-1]),
+                          a.Instrument.xrs & a.goes.SatelliteNumber(16) & a.Resolution("avg1m"))
+file_goes = Fido.fetch(result_goes)
+goes = ts.TimeSeries(file_goes, source='XRS', concatenate=True)
+
+date_format = DateFormatter('%d-%H:%M')
+
+
+fig, full_axs = plt.subplots(3, 2, figsize=(8, 6), gridspec_kw={"width_ratios": [1, 0.03]})
+axs = full_axs[:, 0]
+[ax.axis('off') for ax in full_axs[:, 1]]
+# make date axis
+for ax in axs:
+    ax.xaxis_date()
+    ax.set_xlim(times[0], times[-1])
+    ax.xaxis.set_major_formatter(date_format)
+
+fig.autofmt_xdate()
+
+ax = axs[0]
+ax.plot(times, free_energy)
+ax.set_ylabel('Free Energy\n[$10^{32}$ erg]')
+ax.set_xlim(times[0], times[-1])
+
+ax = axs[1]
+dt = np.diff(filled_times)[0] / 2
+dz = Mm_per_pixel
+max_height = filled_energy_distribution.shape[1] * Mm_per_pixel
+im = ax.imshow(filled_energy_distribution.T,  # average
+                extent=(filled_times[0] - dt, filled_times[-1] + dt, -dz, max_height + dz), aspect='auto',
+                origin='lower',
+                cmap=cm.get_cmap('jet'), norm=LogNorm(vmin=1e17))
+ax.set_ylabel('Height\n[Mm]')
+ax.set_ylim(0, 20)
+# add colorbar
+full_axs[1, 1].axis('on')
+cbar = fig.colorbar(im, cax=full_axs[1, 1], label='Free Energy Density\n' + r'[erg $cm^{-1}]$')
+# cbar.formatter.set_powerlimits((0, 0))
+# cbar.set_ticks([2e2, 4e2, 6e2, 8e2])
+
+ax = axs[2]
+ax.plot(goes.data.index, goes.data['xrsb'], label='xrsb')
+ax.plot(goes.data.index, goes.data['xrsa'], label='xrsa')
+ax.legend( loc='upper left')
+ax.set_yscale("log")
+ax.set_ylim(1e-6, 1e-3)
+ax.set_ylabel('GOES\n' + r'[W m$^{-2}$]')
+ax.set_xlim(times[0], times[-1])
+
+for value in [1e-5, 1e-4]:
+    ax.axhline(value, c='gray')
+
+
+for value, label in zip([1e-6, 1e-5, 1e-4], ['C', 'M', 'X']):
+    ax.text(1.02, value, label, transform=ax.get_yaxis_transform(), horizontalalignment='center')
+
+
+for ax in axs:
+    ax.axvline(x=datetime(2024, 5, 7, 0, 0), linestyle='dotted', c='black')
+    ax.axvline(x=datetime(2024, 5, 7, 20, 18), linestyle='dotted', c='black')
+    ax.axvline(x=datetime(2024, 5, 8, 4, 37), linestyle='dotted', c='black')
+    ax.axvline(x=datetime(2024, 5, 9, 17, 23), linestyle='dotted', c='black')
+    ax.axvline(x=datetime(2024, 5, 10, 6, 27), linestyle='dotted', c='black')
+
+
+# labels = ['A', 'B', 'C', 'M', 'X']
+# centers = np.logspace(-7.5, -3.5, len(labels))
+# for value, label in zip(centers, labels):
+#     ax.text(1.02, value, label, transform=ax.get_yaxis_transform(), horizontalalignment='center')
+
+# f_axs = axs[[0, 1, 3, 4]]
+# if add_flares:
+#     flares = Fido.search(a.Time(min(times), max(times)),
+#                          a.hek.EventType("FL"),
+#                          a.hek.OBS.Observatory == "GOES")["hek"]
+#     goes_mapping = {c: 10 ** (i) for i, c in enumerate(['B', 'C', 'M', 'X'])}
+#     for t, goes_class in zip(flares['event_peaktime'], flares['fl_goescls']):
+#         flare_intensity = np.log10(float(goes_class[1:]) * goes_mapping[goes_class[0]])
+#         if flare_intensity >= np.log10(5 * 10 ** 2):
+#             [ax.axvline(x=date2num(t.datetime), linestyle='dotted', c='black') for ax in f_axs]
+#         if flare_intensity >= np.log10(1 * 10 ** 3):
+#             [ax.axvline(x=date2num(t.datetime), linestyle='dotted', c='red') for ax in f_axs]
+plt.tight_layout()
+plt.savefig(os.path.join(result_path, 'integrated_quantities.jpg'), dpi=300)
+plt.close()
diff --git a/nf2/evaluation/sharp/plot_maps.py b/nf2/evaluation/sharp/plot_maps.py
new file mode 100644
index 0000000..abb9f7b
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_maps.py
@@ -0,0 +1,120 @@
+import argparse
+import glob
+import os
+from copy import copy
+from datetime import datetime
+
+import numpy as np
+from astropy import units as u
+from astropy.io import fits
+from dateutil.parser import parse
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, Normalize
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.sharp.convert_series import load_results
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+    parser.add_argument('--pkl_path', type=str, help='path to the directory with the converted pkl files.')
+    parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+    parser.add_argument('--euv_path', type=str, help='path to the EUV files.', required=False, default=None)
+    args = parser.parse_args()
+
+    output = load_results(args.pkl_path)
+
+    euv_files = np.array(sorted(glob.glob(args.euv_path)))
+    euv_dates = np.array([parse(fits.getheader(f, 1)['DATE-OBS']) for f in euv_files])
+
+    result_path = args.result_path
+    os.makedirs(result_path, exist_ok=True)
+
+    times = output['times']
+    wcs = output['wcs']
+    maps = output['maps']
+    Mm_per_pixel = output['Mm_per_pixel']
+
+    j_norm = LogNorm()
+    free_energy_norm = LogNorm(1e9, 1e14)
+
+    for i, time in tqdm(enumerate(times), total=len(times)):
+        b0_path = os.path.join(result_path, f'b0_{time.isoformat("T", timespec="minutes")}.jpg')
+        euv_path = os.path.join(result_path, f'euv_{time.isoformat("T", timespec="minutes")}.jpg')
+        j_path = os.path.join(result_path, f'j_{time.isoformat("T", timespec="minutes")}.jpg')
+        free_energy_path = os.path.join(result_path, f'free_energy_{time.isoformat("T", timespec="minutes")}.jpg')
+
+        euv_file = euv_files[np.argmin(np.abs(euv_dates - time))]
+        euv_map = Map(euv_file)
+        exposure = euv_map.exposure_time.to_value(u.s)
+        euv_map = euv_map.reproject_to(wcs[i][0])
+
+        b_0 = maps['b_0'][i].value
+        current_density_map = maps['current_density_map'][i].to_value(u.G * u.cm * u.s ** -1)
+        free_energy_map = maps['free_energy_map'][i].to_value(u.erg * u.cm ** -2)
+
+        fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+        extent = np.array([0, b_0.shape[0], 0, b_0.shape[1]]) * Mm_per_pixel
+        b0_im = ax.imshow(b_0.T, origin='lower', cmap='gray', extent=extent, norm=Normalize(vmin=-2000, vmax=2000))
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="3%", pad=0.05)
+        # fig.colorbar(im, cax=cax, label='[G]')
+        plt.tight_layout(pad=0)
+        plt.savefig(b0_path, dpi=300)
+        plt.close()
+
+        fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+        cm = copy(plt.get_cmap('sdoaia131'))
+        cm.set_bad('black')
+        euv_im = ax.imshow(euv_map.data / exposure, origin='lower', cmap=cm, extent=extent, norm=LogNorm(vmin=5, vmax=1e3))
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="3%", pad=0.05)
+        # fig.colorbar(im, cax=cax, label='[DN / s]')
+        plt.tight_layout(pad=0)
+        plt.savefig(euv_path, dpi=300)
+        plt.close()
+
+        fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+        cd_im = ax.imshow(current_density_map.T, origin='lower', cmap='inferno', extent=extent, norm=j_norm)
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="3%", pad=0.05)
+        # fig.colorbar(im, cax=cax, label='[G cm / s]')
+        plt.tight_layout(pad=0)
+        plt.savefig(j_path, dpi=300)
+        plt.close()
+
+        fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+        fe_im = ax.imshow(free_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=free_energy_norm)
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="3%", pad=0.05)
+        # fig.colorbar(im, cax=cax, label='[erg / cm$^2$]')
+        plt.tight_layout(pad=0)
+        plt.savefig(free_energy_path, dpi=300)
+        plt.close()
+
+    fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+    ax.set_axis_off()
+    cbar = fig.colorbar(b0_im, ax=ax, label='[G]')
+    # set ticks
+    cbar.set_ticks([-2000, -1000, 0, 1000, 2000])
+    plt.tight_layout(pad=0)
+    plt.savefig(os.path.join(result_path, f'b0_colorbar.jpg'), dpi=300)
+    plt.close()
+
+    fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+    ax.set_axis_off()
+    fig.colorbar(euv_im, ax=ax, label='[DN / s]')
+    plt.tight_layout(pad=0)
+    plt.savefig(os.path.join(result_path, f'euv_colorbar.jpg'), dpi=300)
+    plt.close()
+
+    fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
+    ax.set_axis_off()
+    fig.colorbar(cd_im, ax=ax, label='[G cm / s]')
+    plt.tight_layout(pad=0)
+    plt.savefig(os.path.join(result_path, f'j_colorbar.jpg'), dpi=300)
+    plt.close()
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/nf2/evaluation/vsm_corrugation.py b/nf2/evaluation/solis/vsm_corrugation.py
similarity index 100%
rename from nf2/evaluation/vsm_corrugation.py
rename to nf2/evaluation/solis/vsm_corrugation.py
diff --git a/nf2/evaluation/spherical/__init__.py b/nf2/evaluation/spherical/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/evaluation/spherical/b_diff.py b/nf2/evaluation/spherical/b_diff.py
new file mode 100644
index 0000000..d4d2366
--- /dev/null
+++ b/nf2/evaluation/spherical/b_diff.py
@@ -0,0 +1,61 @@
+import os
+
+import numpy as np
+import torch
+from astropy import units as u
+from matplotlib import pyplot as plt
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
+from torch import nn
+
+from nf2.data.util import spherical_to_cartesian, vector_cartesian_to_spherical
+from nf2.evaluation.unpack import load_coords
+
+nf2_file = '/glade/work/rjarolim/nf2/global/2154_vp_v11/extrapolation_result.nf2'
+results_path = '/glade/work/rjarolim/nf2/global/2154_vp_v11/results'
+ref_r_file = '/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Br.fits'
+ref_t_file = '/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bt.fits'
+ref_p_file = '/glade/work/rjarolim/data/global/fd_2154/full_disk/hmi.b_720s.20140902_000000_TAI.Bp.fits'
+os.makedirs(results_path, exist_ok=True)
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+ref_r_map = Map(ref_r_file)
+ref_t_map = Map(ref_t_file)
+ref_p_map = Map(ref_p_file)
+coords = all_coordinates_from_map(ref_r_map).transform_to(frames.HeliographicCarrington)
+ref_b = np.stack([ref_r_map.data, -ref_t_map.data, ref_p_map.data]).T
+
+spherical_coords = np.stack([
+    coords.radius.to(u.solRad).value,
+    np.pi / 2 + coords.lat.to(u.rad).value,
+    coords.lon.to(u.rad).value,
+]).transpose()
+cartesian_coords = spherical_to_cartesian(spherical_coords)
+
+# Load the model
+state = torch.load(nf2_file, map_location=device)
+model = nn.DataParallel(state['model'])
+
+b = load_coords(model, 1, state['b_norm'], cartesian_coords, device, progress=True, compute_currents=False)
+
+b = vector_cartesian_to_spherical(b, spherical_coords)
+
+# create difference plot
+fig, axs = plt.subplots(3, 3, figsize=(15, 15))
+for i in range(3):
+    im = axs[i, 0].imshow(ref_b[..., i].T, cmap='grey', vmin=-500, vmax=500)
+    plt.colorbar(im, ax=axs[i, 0], pad=0.1, shrink=0.8, label='B [G]')
+    axs[i, 0].set_title('Reference')
+    im = axs[i, 1].imshow(b[..., i].T, cmap='grey', vmin=-500, vmax=500)
+    plt.colorbar(im, ax=axs[i, 1], pad=0.1, shrink=0.8, label='B [G]')
+    axs[i, 1].set_title('NF2')
+    im = axs[i, 2].imshow(np.abs(b[..., i] - ref_b[..., i]).T, cmap='viridis', vmin=0, vmax=100)
+    plt.colorbar(im, ax=axs[i, 2], pad=0.1, shrink=0.8, label='Difference [G]')
+    axs[i, 2].set_title('Difference')
+plt.tight_layout()
+plt.savefig(os.path.join(results_path, f'difference.jpg'), dpi=300)
+plt.close()
+
+
+print('Mean absolute difference:', np.nanmean(((b - ref_b) ** 2).sum(-1) ** 0.5))
diff --git a/nf2/evaluation/spherical/b_diff_synoptic.py b/nf2/evaluation/spherical/b_diff_synoptic.py
new file mode 100644
index 0000000..f890a5b
--- /dev/null
+++ b/nf2/evaluation/spherical/b_diff_synoptic.py
@@ -0,0 +1,81 @@
+import argparse
+import os
+
+import numpy as np
+import torch
+from astropy import units as u
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
+from torch import nn
+
+from nf2.data.util import spherical_to_cartesian, vector_cartesian_to_spherical
+from nf2.evaluation.unpack import load_coords
+
+parser = argparse.ArgumentParser(description='Evaluate NF2 on synoptic data')
+parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+parser.add_argument('--out_path', type=str, help='path to the target VTK files')
+parser.add_argument('--cr', type=int, help='Carrington rotation number')
+
+args = parser.parse_args()
+
+nf2_file = args.nf2_path
+results_path = args.out_path
+carrington_rotation = args.cr
+os.makedirs(results_path, exist_ok=True)
+
+ref_r_file = f"/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.{carrington_rotation}.Mr_polfil.fits",
+ref_t_file = f"/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{carrington_rotation}.Bt.fits"
+ref_p_file = f"/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{carrington_rotation}.Bp.fits",
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+ref_r_map = Map(ref_r_file)
+ref_t_map = Map(ref_t_file)
+ref_p_map = Map(ref_p_file)
+coords = all_coordinates_from_map(ref_r_map).transform_to(frames.HeliographicCarrington)
+ref_b = np.stack([ref_r_map.data, -ref_t_map.data, ref_p_map.data]).T * u.G
+
+spherical_coords = np.stack([
+    coords.radius.value,
+    np.pi / 2 + coords.lat.to(u.rad).value,
+    coords.lon.to(u.rad).value,
+]).transpose()
+cartesian_coords = spherical_to_cartesian(spherical_coords)
+
+# Load the model
+state = torch.load(nf2_file, map_location=device)
+model = nn.DataParallel(state['model'])
+
+b = load_coords(model, 1, state['data']['G_per_dB'], cartesian_coords, device, progress=True, compute_currents=False)
+
+b = vector_cartesian_to_spherical(b, spherical_coords)
+
+extent = [coords.lon.min().value, coords.lon.max().value, coords.lat.min().value, coords.lat.max().value]
+# create difference plot
+fig, axs = plt.subplots(3, 3, figsize=(25, 15))
+for i in range(3):
+    im = axs[i, 0].imshow(ref_b[..., i].value.T, cmap='grey', vmin=-500, vmax=500, origin='lower', extent=extent)
+    divider = make_axes_locatable(axs[i, 0])
+    cax = divider.append_axes("right", size="2%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='B$_r$ [G]')
+
+    im = axs[i, 1].imshow(b[..., i].value.T, cmap='grey', vmin=-500, vmax=500, origin='lower', extent=extent)
+    divider = make_axes_locatable(axs[i, 1])
+    cax = divider.append_axes("right", size="2%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='B$_{\\theta}$ [G]')
+
+    im = axs[i, 2].imshow(np.abs(b[..., i] - ref_b[..., i]).value.T, cmap='viridis', vmin=0, vmax=100, origin='lower', extent=extent)
+    divider = make_axes_locatable(axs[i, 2])
+    cax = divider.append_axes("right", size="2%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='B$_{\\phi}$ [G]')
+
+axs[0, 0].set_title('Reference')
+axs[0, 1].set_title('NF2')
+axs[0, 2].set_title('Difference')
+plt.tight_layout()
+plt.savefig(os.path.join(results_path, f'difference.jpg'), dpi=300)
+plt.close()
+
+print('Mean absolute difference:', np.nanmean(((b - ref_b) ** 2).sum(-1) ** 0.5))
diff --git a/nf2/evaluation/spherical/compare.py b/nf2/evaluation/spherical/compare.py
new file mode 100644
index 0000000..ffc4745
--- /dev/null
+++ b/nf2/evaluation/spherical/compare.py
@@ -0,0 +1,314 @@
+import numpy as np
+from astropy import units as u
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
+
+from nf2.data.util import vector_cartesian_to_spherical
+from nf2.evaluation.metric import theta_J, divergence_jacobian, energy, b_diff_error
+from nf2.evaluation.output import CartesianOutput, SphericalOutput, current_density, spherical_energy_gradient, \
+    energy_gradient
+from nf2.potential.potential_field import get_potential_field
+
+Mm_per_pixel = 0.72
+
+cartesian_model = CartesianOutput('/glade/work/rjarolim/nf2/sharp/377_v01/extrapolation_result.nf2')
+cartesian_out = cartesian_model.load_cube(progress=True, Mm_per_pixel=Mm_per_pixel, metrics={'j' : current_density, 'energy_gradient': energy_gradient})
+
+reference_map = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br.fits')
+
+bottom_left = reference_map.bottom_left_coord.transform_to(frames.HeliographicCarrington)
+top_right = reference_map.top_right_coord.transform_to(frames.HeliographicCarrington)
+center = reference_map.center.transform_to(frames.HeliographicCarrington)
+
+spherical_model = SphericalOutput('/glade/work/rjarolim/nf2/spherical/377_v05/extrapolation_result.nf2')
+spherical_out = spherical_model.load_spherical(
+    latitude_range=((np.pi / 2 * u.rad) - top_right.lat, (np.pi / 2 * u.rad) - bottom_left.lat),
+    longitude_range=(bottom_left.lon, top_right.lon),
+    radius_range=(1. * u.solRad, 1 * u.solRad + 100 * u.Mm),
+    sampling=[cartesian_out['coords'].shape[2], cartesian_out['coords'].shape[1], cartesian_out['coords'].shape[0]],
+    metrics={'j' : current_density, 'energy_gradient': spherical_energy_gradient},
+    progress=True)
+
+subframe_model = SphericalOutput('/glade/work/rjarolim/nf2/spherical/377_subframe_v04/extrapolation_result.nf2')
+subframe_out = subframe_model.load_spherical(
+    latitude_range=((np.pi / 2 * u.rad) - top_right.lat, (np.pi / 2 * u.rad) - bottom_left.lat),
+    longitude_range=(bottom_left.lon, top_right.lon),
+    radius_range=(1. * u.solRad, 1 * u.solRad + 100 * u.Mm),
+    sampling=[cartesian_out['coords'].shape[2], cartesian_out['coords'].shape[1], cartesian_out['coords'].shape[0]],
+    metrics={'j' : current_density, 'energy_gradient': spherical_energy_gradient},
+    progress=True)
+
+pb_model = SphericalOutput('/glade/work/rjarolim/nf2/spherical/377_potential_v01/extrapolation_result.nf2')
+pb_out = spherical_model.load_spherical(
+    latitude_range=((np.pi / 2 * u.rad) - top_right.lat, (np.pi / 2 * u.rad) - bottom_left.lat),
+    longitude_range=(bottom_left.lon, top_right.lon),
+    radius_range=(1. * u.solRad, 1 * u.solRad + 100 * u.Mm),
+    sampling=[cartesian_out['coords'].shape[2], cartesian_out['coords'].shape[1], cartesian_out['coords'].shape[0]],
+    metrics={'j' : current_density, 'energy_gradient': spherical_energy_gradient},
+    progress=True)
+
+b_pb = pb_out['b'].value
+b_spherical = spherical_out['b'].value
+b_cartesian = cartesian_out['b'].value
+b_subframe = subframe_out['b'].value
+
+j_pb = pb_out['j'].value
+j_spherical = spherical_out['j'].value
+j_cartesian = cartesian_out['j'].value
+j_subframe = subframe_out['j'].value
+
+energy_gradient_pb = pb_out['energy_gradient'].value
+energy_gradient_spherical = spherical_out['energy_gradient'].value
+energy_gradient_cartesian = cartesian_out['energy_gradient'].value
+energy_gradient_subframe = subframe_out['energy_gradient'].value
+
+coords_pb = pb_out['coords']
+coords_spherical = spherical_out['coords']
+coords_cartesian = cartesian_out['coords']
+coords_subframe = subframe_out['coords']
+
+div_pb = np.nanmean(np.abs(divergence_jacobian(pb_out['jac_matrix'].value)))
+div_spherical = np.nanmean(np.abs(divergence_jacobian(spherical_out['jac_matrix'].value)))
+div_cartesian = np.nanmean(np.abs(divergence_jacobian(cartesian_out['jac_matrix'].value)))
+div_subframe = np.nanmean(np.abs(divergence_jacobian(subframe_out['jac_matrix'].value)))
+
+theta_pb = theta_J(b_pb, j_pb)
+theta_spherical = theta_J(b_spherical, j_spherical)
+theta_cartesian = theta_J(b_cartesian, j_cartesian)
+theta_subframe = theta_J(b_subframe, j_subframe)
+
+bx_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp.fits')
+by_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt.fits')
+bz_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br.fits')
+b_cartesian_ref = np.stack([bx_ref.data, -by_ref.data, bz_ref.data]).T
+
+errx_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bp_err.fits')
+erry_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Bt_err.fits')
+errz_ref = Map('/glade/work/rjarolim/data/nf2/377/hmi.sharp_cea_720s.377.20110215_000000_TAI.Br_err.fits')
+err_cartesian_ref = np.stack([errx_ref.data, erry_ref.data, errz_ref.data]).T
+
+spherical_coords = all_coordinates_from_map(reference_map).T
+#
+img_out = cartesian_model.load_boundary()
+b_cartesian_img = img_out['b'].value
+#
+img_out = spherical_model.load_spherical_coords(spherical_coords)
+b_spherical_img = vector_cartesian_to_spherical(img_out['b'].value, img_out['spherical_coords'])
+b_spherical_img = np.flip(b_spherical_img, -1)
+b_spherical_img[..., 1] *= -1
+#
+img_out = subframe_model.load_spherical_coords(spherical_coords)
+b_subframe_img = vector_cartesian_to_spherical(img_out['b'].value, img_out['spherical_coords'])
+b_subframe_img = np.flip(b_subframe_img, -1)
+b_subframe_img[..., 1] *= -1
+#
+img_out = pb_model.load_spherical_coords(spherical_coords)
+b_pb_img = vector_cartesian_to_spherical(img_out['b'].value, img_out['spherical_coords'])
+b_pb_img = np.flip(b_pb_img, -1)
+b_pb_img[..., 1] *= -1
+
+b_diff_cartesian = np.linalg.norm(b_cartesian_img - b_cartesian_ref, axis=-1).mean()
+b_diff_spherical = np.linalg.norm(b_spherical_img - b_cartesian_ref, axis=-1).mean()
+b_diff_subframe = np.linalg.norm(b_subframe_img - b_cartesian_ref, axis=-1).mean()
+b_dff_pb = np.linalg.norm(b_pb_img - b_cartesian_ref, axis=-1).mean()
+
+b_diff_err_cartesian = b_diff_error(b_cartesian_img, b_cartesian_ref, err_cartesian_ref).mean()
+b_diff_err_spherical = b_diff_error(b_spherical_img, b_cartesian_ref, err_cartesian_ref).mean()
+b_diff_err_subframe = b_diff_error(b_subframe_img, b_cartesian_ref, err_cartesian_ref).mean()
+b_diff_err_pb = b_diff_error(b_pb_img, b_cartesian_ref, err_cartesian_ref).mean()
+
+print(f'Divergence: Cartesian: {div_cartesian:.2e}, Spherical:  {div_spherical:.2e}, Subframe: {div_subframe:.2e}, PB: {div_pb:.2e}')
+print(f'Weighted Theta:  Cartesian: {theta_cartesian:.2f}, Spherical: {theta_spherical:.2f}, Subframe: {theta_subframe:.2f}, PB: {theta_pb:.2f}')
+print(f'B Diff: Cartesian: {b_diff_cartesian:.2f}, Spherical: {b_diff_spherical:.2f}, Subframe: {b_diff_subframe:.2f}, PB: {b_dff_pb:.2f}')
+print(f'B Diff Err: Cartesian: {b_diff_err_cartesian:.2f}, Spherical: {b_diff_err_spherical:.2f}, Subframe: {b_diff_err_subframe:.2f}, PB: {b_diff_err_pb:.2f}')
+
+# print(f'Saving Cube (spherical): {coords_spherical.shape}')
+# save_vtk('/glade/work/rjarolim/nf2/spherical/evaluation/spherical.vtk', coords_spherical, vectors={'b': np.nan_to_num(b_spherical), 'j': np.nan_to_num(j_spherical)}, Mm_per_pix=0.72)
+# print(f'Saving Cube (cartesian): {coords_cartesian.shape}')
+# save_vtk('/glade/work/rjarolim/nf2/spherical/evaluation/cartesian.vtk', coords_cartesian, vectors={'b': b_cartesian, 'j': j_cartesian}, Mm_per_pix=0.72)
+
+
+b_potential = get_potential_field(b_cartesian[:, :, 0, 2], b_cartesian.shape[2], batch_size=int(1e3))
+cartesian_free_energy = energy(b_cartesian) - energy(b_potential)
+sphere_free_energy = energy(b_spherical) - np.moveaxis(energy(b_potential), (0, 1, 2), (2, 1, 0))
+subframe_free_energy = energy(b_subframe) - np.moveaxis(energy(b_potential), (0, 1, 2), (2, 1, 0))
+
+# plot comparison of currents and jxb
+
+cartesian_extent = [0, b_cartesian.shape[0] * Mm_per_pixel, 0, b_cartesian.shape[1] * Mm_per_pixel]
+spherical_extent = [bottom_left.lon.to(u.deg).value, top_right.lon.to(u.deg).value,
+          top_right.lat.to(u.deg).value, bottom_left.lat.to(u.deg).value, ]
+
+fig, axs = plt.subplots(3, 3, figsize=(12, 8), )
+
+ax = axs[0, 0]
+im = ax.imshow(b_cartesian_img[..., 2].T, origin='lower', cmap='gray', extent=cartesian_extent, vmin=-1000, vmax=1000)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+ax.set_title('Cartesian')
+ax.set_ylabel('Y [Mm]')
+ax.set_xlabel('X [Mm]')
+
+ax = axs[0, 1]
+im = ax.imshow(b_spherical_img[..., 2].T, origin='lower', cmap='gray', extent=spherical_extent, vmin=-1000, vmax=1000)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+ax.set_title('Spherical')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+ax = axs[0, 2]
+im = ax.imshow(b_subframe_img[..., 2].T, origin='lower', cmap='gray', extent=spherical_extent, vmin=-1000, vmax=1000)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+ax.set_title('Subframe')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+norm = LogNorm()
+ax = axs[1, 0]
+im = ax.imshow(np.linalg.norm(j_cartesian, axis=-1).sum(2).T * Mm_per_pixel * 1e6, origin='lower',
+                 cmap='inferno', norm=norm, extent=cartesian_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'|$\nabla x $B| [G])')
+ax.set_title('Cartesian')
+ax.set_ylabel('Y [Mm]')
+ax.set_xlabel('X [Mm]')
+
+ax = axs[1, 1]
+im = ax.imshow(np.nansum(np.linalg.norm(j_spherical, axis=-1), 0) * Mm_per_pixel * 1e6, origin='upper',
+                 cmap='inferno', norm=norm, extent=spherical_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'|$\nabla x $B| [G])')
+ax.set_title('Spherical')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+ax = axs[1, 2]
+im = ax.imshow(np.nansum(np.linalg.norm(j_subframe, axis=-1), 0) * Mm_per_pixel * 1e6, origin='upper',
+                    cmap='inferno', norm=norm, extent=spherical_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'|$\nabla x $B| [G])')
+ax.set_title('Subframe')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+
+norm = LogNorm(vmin=1e9)
+ax = axs[2, 0]
+im = ax.imshow(cartesian_free_energy.sum(2).T * Mm_per_pixel * 1e6, origin='lower', cmap='viridis', norm=norm,
+                    extent=cartesian_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Energy [erg/cm$^3$]')
+ax.set_title('Cartesian')
+ax.set_ylabel('Y [Mm]')
+ax.set_xlabel('X [Mm]')
+
+ax = axs[2, 1]
+im = ax.imshow(sphere_free_energy.sum(0) * Mm_per_pixel * 1e6, origin='upper', cmap='viridis', norm=norm,
+                    extent=spherical_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Energy [erg/cm$^3$]')
+ax.set_title('Spherical')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+ax = axs[2, 2]
+im = ax.imshow(subframe_free_energy.sum(0) * Mm_per_pixel * 1e6, origin='upper', cmap='viridis', norm=norm,
+                    extent=spherical_extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Energy [erg/cm$^3$]')
+ax.set_title('Subframe')
+ax.set_ylabel('Latitude [deg]')
+ax.set_xlabel('Longitude [deg]')
+
+fig.tight_layout()
+plt.savefig('/glade/work/rjarolim/nf2/spherical/evaluation/overview.png', dpi=300, transparent=True)
+plt.close()
+
+
+
+
+fig, axs = plt.subplots(1, 5, figsize=(15, 3), )
+
+ax = axs[0]
+cartesian_energy = energy(b_cartesian).sum((0, 1)) * Mm_per_pixel ** 2
+spherical_energy = energy(b_spherical).sum((1, 2)) * Mm_per_pixel ** 2
+subframe_energy = energy(b_subframe).sum((1, 2)) * Mm_per_pixel ** 2
+pb_energy = energy(b_pb).sum((1, 2)) * Mm_per_pixel ** 2
+height = np.linspace(0, b_cartesian.shape[2] * Mm_per_pixel, b_cartesian.shape[2])
+ax.plot(cartesian_energy, height, label='Cartesian')
+ax.plot(spherical_energy, height, label='Spherical')
+ax.plot(subframe_energy, height, label='Subframe')
+ax.plot(pb_energy, height, label='Potential Boundary')
+ax.set_ylabel('Height [Mm]')
+ax.set_xlabel('Energy [erg/Mm]')
+ax.set_ylim(0, 100)
+ax.semilogx()
+
+# currents
+ax = axs[1]
+ax.plot(np.linalg.norm(j_cartesian, axis=-1).mean((0, 1)), height, label='Cartesian')
+ax.plot(np.linalg.norm(j_spherical, axis=-1).mean((1, 2)), height, label='Spherical')
+ax.plot(np.linalg.norm(j_subframe, axis=-1).mean((1, 2)), height, label='Subframe')
+ax.plot(np.linalg.norm(j_pb, axis=-1).mean((1, 2)), height, label='Potential Boundary')
+ax.set_ylabel('Height [Mm]')
+ax.set_xlabel('|J| [G/Mm]')
+ax.set_ylim(0, 100)
+ax.semilogx()
+
+ax = axs[2]
+cartesian_free_energy_h = cartesian_free_energy.sum((0, 1)) * Mm_per_pixel ** 2
+spherical_free_energy_h = sphere_free_energy.sum((1, 2)) * Mm_per_pixel ** 2
+subframe_free_energy_h = subframe_free_energy.sum((1, 2)) * Mm_per_pixel ** 2
+ax.plot(cartesian_free_energy_h, height, label='Cartesian')
+ax.plot(spherical_free_energy_h, height, label='Spherical')
+ax.plot(subframe_free_energy_h, height, label='Subframe')
+ax.plot(pb_energy, height, label='Potential Boundary')
+ax.set_ylabel('Height [Mm]')
+ax.set_xlabel('Free Energy [erg/Mm]')
+ax.set_ylim(0, 100)
+ax.semilogx()
+
+ax = axs[3]
+cartesian_jxb = np.linalg.norm(np.cross(j_cartesian, b_cartesian), axis=-1).mean((0, 1))
+spherical_jxb = np.linalg.norm(np.cross(j_spherical, b_spherical), axis=-1).mean((1, 2))
+subframe_jxb = np.linalg.norm(np.cross(j_subframe, b_subframe), axis=-1).mean((1, 2))
+pb_jxb = np.linalg.norm(np.cross(j_pb, b_pb), axis=-1).mean((1, 2))
+ax.plot(cartesian_jxb, height, label='Cartesian')
+ax.plot(spherical_jxb, height, label='Spherical')
+ax.plot(subframe_jxb, height, label='Subframe')
+ax.plot(pb_jxb, height, label='Potential Boundary', linestyle='--')
+ax.set_ylabel('Height [Mm]')
+ax.set_xlabel('|JxB| [G$^2$/Mm]')
+ax.set_ylim(0, 100)
+ax.semilogx()
+
+
+ax = axs[4]
+ax.plot(energy_gradient_cartesian.sum((0, 1)), height, label='Cartesian')
+ax.plot(energy_gradient_spherical.sum((1, 2)), height, label='Spherical')
+ax.plot(energy_gradient_subframe.sum((1, 2)), height, label='Subframe')
+ax.plot(energy_gradient_pb.sum((1, 2)), height, label='Potential Boundary')
+ax.set_ylabel('Height [Mm]')
+ax.set_xlabel('Energy Gradient [erg/Mm]')
+ax.set_ylim(0, 100)
+ax.set_xlim(-2, 10)
+ax.axvline(0, color='k', linestyle='--')
+
+axs[-1].legend(fontsize=16)
+fig.tight_layout()
+plt.savefig('/glade/work/rjarolim/nf2/spherical/evaluation/height.png', dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/spherical/current_change.py b/nf2/evaluation/spherical/current_change.py
new file mode 100644
index 0000000..87cbc39
--- /dev/null
+++ b/nf2/evaluation/spherical/current_change.py
@@ -0,0 +1,120 @@
+import glob
+import os
+
+import numpy as np
+from astropy.coordinates import SkyCoord
+from astropy import units as u
+from matplotlib.colors import SymLogNorm
+from sunpy.coordinates import frames
+from sunpy.map import Map, make_heliographic_header, all_coordinates_from_map
+
+from nf2.data.util import vector_cartesian_to_spherical
+from nf2.evaluation.output import SphericalOutput
+
+import matplotlib.pyplot as plt
+
+result_path = '/glade/work/rjarolim/nf2/spherical/2283_00_v01'
+pre_eruption = SphericalOutput('/glade/work/rjarolim/nf2/spherical/2283_00_v01/extrapolation_result.nf2')
+post_eruption = SphericalOutput('/glade/work/rjarolim/nf2/spherical/2283_04_v01/extrapolation_result.nf2')
+
+latitude_range = np.deg2rad(-45), np.deg2rad(45)
+longitude_range = 1.673, 4.501
+
+coords = np.stack(np.meshgrid(
+np.linspace(1.0, 1.3, 128),
+    np.linspace(*latitude_range, np.rad2deg(latitude_range[1] - latitude_range[0]).astype(int) * 2),
+    np.linspace(*longitude_range, np.rad2deg(longitude_range[1] - longitude_range[0]).astype(int) * 2), indexing='ij'),
+    -1)
+
+# PFSS extrapolation
+spherical_boundary_coords = SkyCoord(lon=coords[..., 2] * u.rad, lat=(coords[..., 1]) * u.rad,
+                                     radius=coords[..., 0] * u.solRad, frame=frames.HeliographicCarrington)
+
+pre_eruption_out = pre_eruption.load_spherical_coords(spherical_boundary_coords, progress=True, batch_size=4096)
+post_eruption_out = post_eruption.load_spherical_coords(spherical_boundary_coords, progress=True, batch_size=4096)
+
+pre_j = pre_eruption_out['j'].to_value(u.G / u.Mm)
+pre_j = vector_cartesian_to_spherical(pre_j, pre_eruption_out['spherical_coords'])[..., 1:]
+post_j = post_eruption_out['j'].to_value(u.G / u.Mm)
+post_j = vector_cartesian_to_spherical(post_j, post_eruption_out['spherical_coords'])[..., 1:]
+
+pre_current_density = np.linalg.norm(pre_j, axis=-1)
+post_current_density = np.linalg.norm(post_j, axis=-1)
+
+current_change = post_current_density - pre_current_density
+current_change = np.clip(current_change, a_min=None, a_max=0)
+
+# translate AIA maps to spherical coordinates
+shape = (720, 1440)
+aia_maps = sorted(glob.glob('/glade/work/rjarolim/data/global/fd_2283/aia.lev1_euv_12s.*.131.image_lev1.fits'))
+
+aia_data = []
+
+for f in aia_maps:
+    aia_map = Map(f)
+    carr_header = make_heliographic_header(aia_map.date, aia_map.observer_coordinate, shape, frame='carrington')
+    aia_map = aia_map.reproject_to(carr_header)
+    d = aia_map.data
+    d = np.roll(d, 720, axis=1)
+    aia_data.append(d)
+
+pre_aia = aia_data[0]
+post_aia = aia_data[-1]
+aia_integrated = np.mean(aia_data, axis=0)
+
+
+# plot current change maps
+extent = [np.rad2deg(longitude_range[0]), np.rad2deg(longitude_range[1]),
+            np.rad2deg(latitude_range[0]), np.rad2deg(latitude_range[1])]
+
+fig, axs = plt.subplots(2, 3, figsize=(15, 5))
+
+ax = axs[0, 0]
+im = ax.imshow(pre_aia, origin='lower', cmap='sdoaia131',
+               extent=[0, 360, -90, 90], norm=SymLogNorm(linthresh=10, vmin=1, vmax=2e3))
+plt.colorbar(im, ax=ax, label='Intensity [DN/s]')
+ax.set_title('Pre-Eruption')
+ax.set_xlim(*extent[:2])
+ax.set_ylim(*extent[2:])
+
+ax = axs[0, 1]
+im = ax.imshow(post_aia, origin='lower', cmap='sdoaia131',
+               extent=[0, 360, -90, 90], norm=SymLogNorm(linthresh=10, vmin=1, vmax=2e3))
+plt.colorbar(im, ax=ax, label='Intensity [DN/s]')
+ax.set_title('Post-Eruption')
+ax.set_xlim(*extent[:2])
+ax.set_ylim(*extent[2:])
+
+ax = axs[0, 2]
+im = ax.imshow(aia_integrated, origin='lower', cmap='RdBu_r',
+               extent=[0, 360, -90, 90], norm=SymLogNorm(linthresh=10, vmin=-2e3, vmax=2e3))
+plt.colorbar(im, ax=ax, label='Intensity Change [DN/s]')
+ax.set_title('Integrated Intensity')
+ax.set_xlim(*extent[:2])
+ax.set_ylim(*extent[2:])
+
+ax = axs[1, 0]
+im = ax.imshow(pre_current_density.sum(0), origin='lower', cmap='inferno',
+               norm=SymLogNorm(linthresh=10, vmin=0, vmax=1000), extent=extent)
+plt.colorbar(im, ax=ax, label='Current Density [G/Mm]')
+ax.set_title('Pre-Eruption')
+
+ax = axs[1, 1]
+im = ax.imshow(post_current_density.sum(0), origin='lower', cmap='inferno',
+               norm=SymLogNorm(linthresh=10, vmin=0, vmax=1000), extent=extent)
+plt.colorbar(im, ax=ax, label='Current Density [G/Mm]')
+ax.set_title('Post-Eruption')
+
+ax = axs[1, 2]
+im = ax.imshow(current_change.sum(0), origin='lower',
+               norm=SymLogNorm(linthresh=100, vmin=-1000, vmax=1000),
+                cmap='RdBu_r', extent=extent)
+plt.colorbar(im, ax=ax, label='Current Density Change [G/Mm]')
+
+[ax.set_xlabel('Longitude [deg]') for ax in np.ravel(axs)]
+axs[0, 0].set_ylabel('Latitude [deg]')
+axs[1, 0].set_ylabel('Latitude [deg]')
+
+fig.tight_layout()
+plt.savefig(os.path.join(result_path, 'current_change.png'), dpi=300, transparent=True)
+plt.close()
\ No newline at end of file
diff --git a/nf2/evaluation/spherical/currents.py b/nf2/evaluation/spherical/currents.py
new file mode 100644
index 0000000..ade32c0
--- /dev/null
+++ b/nf2/evaluation/spherical/currents.py
@@ -0,0 +1,146 @@
+import glob
+import os
+
+import numpy as np
+import pfsspy
+import torch
+from astropy.coordinates import SkyCoord
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, SymLogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from torch import nn
+
+from nf2.data.util import spherical_to_cartesian, vector_spherical_to_cartesian
+from nf2.evaluation.unpack import load_coords
+
+from astropy import units as u
+
+# cr = 2213
+
+files = f'/glade/work/rjarolim/nf2/global/2173_vp_series_v1/*.nf2'
+results_path = '/glade/work/rjarolim/nf2/global/2173_vp_series_v1/results'
+synoptic_map_path = f'/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits'
+
+os.makedirs(results_path, exist_ok=True)
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+files = sorted(glob.glob(files))[::5]
+
+for file in files:
+    state = torch.load(file, map_location=device)
+    model = nn.DataParallel(state['model'])
+
+    radius = 1.3
+    coords = np.stack(np.meshgrid(
+                            np.linspace(1, radius, 128),
+                            np.linspace(0, np.pi, 180),
+                            np.linspace(0, 2 * np.pi, 360), indexing='ij'), -1)
+
+    cartesian_coords = spherical_to_cartesian(coords)
+
+    b, j = load_coords(model, 1, state['b_norm'], cartesian_coords, device, progress=True, compute_currents=True)
+
+    extent = [0, 360, -90, 90]
+
+    # currents
+    fig, ax = plt.subplots(figsize=(10, 5))
+    current_density = (j ** 2).sum(-1) ** 0.5
+    im = ax.imshow(current_density.sum(0), origin='lower',
+                   norm=LogNorm(vmin=1e3, vmax=1e6),
+                   cmap='inferno',
+                   extent=extent)
+    ax.set_title('Currents')
+    ax.set_xlabel('Longitude [deg]')
+    ax.set_ylabel('Latitude [deg]')
+    ax.grid()
+    ax.set_xticks(np.arange(0, 360, 20))
+    ax.set_yticks(np.arange(-80, 80, 20))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    plt.savefig(os.path.join(results_path, f'currents_{os.path.basename(file).replace("nf2", "jpg")}'),
+                dpi=300)
+    plt.close(fig)
+
+    # plot photospheric magnetic field
+    fig, ax = plt.subplots(figsize=(10, 5))
+    b_rtp = vector_spherical_to_cartesian(b, coords)
+    im = ax.imshow(b[0, :, :, 0], origin='lower',
+                   vmin=-500, vmax=500,
+                   cmap='gray', extent=extent)
+    ax.set_title('Photospheric Magnetic Field')
+    ax.set_xlabel('Longitude [deg]')
+    ax.set_ylabel('Latitude [deg]')
+    ax.grid()
+    ax.set_xticks(np.arange(0, 360, 20))
+    ax.set_yticks(np.arange(-80, 80, 20))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    plt.savefig(os.path.join(results_path, f'b_{os.path.basename(file).replace("nf2", "jpg")}'),
+                dpi=300)
+    plt.close(fig)
+
+    # plot integrated magnetic field
+    fig, ax = plt.subplots(figsize=(10, 5))
+    energy = (b ** 2).sum(-1) ** 0.5
+    im = ax.imshow(energy.sum(0), origin='lower',
+                   norm=LogNorm(vmin=1e1, vmax=1e4),
+                   cmap='viridis', extent=extent)
+    ax.set_title('Integrated Magnetic Field')
+    ax.set_xlabel('Longitude [deg]')
+    ax.set_ylabel('Latitude [deg]')
+    ax.grid()
+    ax.set_xticks(np.arange(0, 360, 20))
+    ax.set_yticks(np.arange(-80, 80, 20))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    plt.savefig(os.path.join(results_path, f'energy_{os.path.basename(file).replace("nf2", "jpg")}'),
+                dpi=300)
+    plt.close(fig)
+
+    # vorticity
+    fig, ax = plt.subplots(figsize=(10, 5))
+    vorticity = np.linalg.norm(j, axis=-1) / np.linalg.norm(b, axis=-1)
+    vorticity[np.linalg.norm(b, axis=-1) < 1] = 0
+    im = ax.imshow(vorticity.sum(0), origin='lower',
+                   vmin=0, vmax=8000,
+                   cmap='inferno', extent=extent)
+    ax.set_title('Vorticity')
+    ax.set_xlabel('Longitude [deg]')
+    ax.set_ylabel('Latitude [deg]')
+    ax.grid()
+    ax.set_xticks(np.arange(0, 360, 20))
+    ax.set_yticks(np.arange(-80, 80, 20))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    plt.savefig(os.path.join(results_path,
+                             f'vorticity_{os.path.basename(file).replace("nf2", "jpg")}'),
+                dpi=300)
+    plt.close(fig)
+
+    # plot lorentz force
+    fig, ax = plt.subplots(figsize=(10, 5))
+    lorentz = np.linalg.norm(np.cross(j, b), axis=-1)
+    im = ax.imshow(lorentz.sum(0), origin='lower',
+                   norm=LogNorm(vmin=1e3, vmax=1e7),
+                   cmap='inferno', extent=extent)
+    ax.set_title('Lorentz Force')
+    ax.set_xlabel('Longitude [deg]')
+    ax.set_ylabel('Latitude [deg]')
+    ax.grid()
+    ax.set_xticks(np.arange(0, 360, 20))
+    ax.set_yticks(np.arange(-80, 80, 20))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    plt.savefig(os.path.join(results_path, f'lorentz_{os.path.basename(file).replace("nf2", "jpg")}'),
+                dpi=300)
+    plt.close(fig)
+
+
+
diff --git a/nf2/evaluation/spherical/energy_gradient.py b/nf2/evaluation/spherical/energy_gradient.py
new file mode 100644
index 0000000..a658965
--- /dev/null
+++ b/nf2/evaluation/spherical/energy_gradient.py
@@ -0,0 +1,74 @@
+import glob
+import os
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, SymLogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from torch import nn
+
+from nf2.data.util import spherical_to_cartesian
+from nf2.evaluation.unpack import load_coords
+
+files = '/glade/work/rjarolim/global/2154_vp_v4'
+results_path = '/glade/work/rjarolim/global/2154_vp_v4/results'
+os.makedirs(results_path, exist_ok=True)
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+files = sorted(glob.glob(os.path.join(files, '*.nf2')))
+
+for file in files:
+    state = torch.load(file, map_location=device)
+    model = nn.DataParallel(state['model'])
+
+    radius = 1.3
+    coords = np.stack(np.meshgrid(
+                            np.linspace(1, radius, 512),
+                            np.linspace(0, 2 * np.pi, 1024),
+                            np.deg2rad(270), indexing='ij'), -1)
+    coords = coords[:, :, 0]
+
+    cartesian_coords = spherical_to_cartesian(coords)
+
+    b, j = load_coords(model, 1, state['b_norm'], cartesian_coords, device, progress=True, compute_currents=True)
+
+    # create polar plot
+    energy = np.linalg.norm(b, axis=-1)
+    fig = plt.figure(figsize=(10, 10))
+    ax = fig.add_subplot(111, projection='polar')
+    thetas_radians, radii = coords[:, :, 1], coords[:, :, 0]
+    im = ax.pcolormesh(thetas_radians, radii, energy, norm=LogNorm(1e1, 1e3), cmap='viridis')
+    ax.set_rlim(0, radius)
+    ax.set_thetalim(0, np.pi)
+    ax.set_theta_zero_location('N')
+    plt.colorbar(im, ax=ax, pad=0.1, shrink=0.8, label='Energy density [ergs/cm$^3$]')
+    plt.savefig(os.path.join(results_path, f'energy_polar_{os.path.basename(file)}.jpg'), dpi=300)
+    plt.close()
+
+    # create current plot
+    current_density = np.linalg.norm(j, axis=-1)
+    fig = plt.figure(figsize=(10, 10))
+    ax = fig.add_subplot(111, projection='polar')
+    thetas_radians, radii = coords[:, :, 1], coords[:, :, 0]
+    im = ax.pcolormesh(thetas_radians, radii, current_density, norm=LogNorm(1e3, 1e5), cmap='plasma')
+    ax.set_rlim(0, radius)
+    ax.set_thetalim(0, np.pi)
+    ax.set_theta_zero_location('N')
+    plt.colorbar(im, ax=ax, pad=0.1, shrink=0.8, label='Current density [G/cm$^3$]')
+    plt.savefig(os.path.join(results_path, f'current_density_polar_{os.path.basename(file)}.jpg'), dpi=300)
+    plt.close()
+
+    # create energy gradient plot
+    energy_gradient = np.gradient(energy, axis=0)
+    fig = plt.figure(figsize=(10, 10))
+    ax = fig.add_subplot(111, projection='polar')
+    thetas_radians, radii = coords[:, :, 1], coords[:, :, 0]
+    im = ax.pcolormesh(thetas_radians, radii, energy_gradient, norm=SymLogNorm(1, vmin=-10, vmax=10), cmap='RdBu_r')
+    ax.set_rlim(0, radius)
+    ax.set_thetalim(0, np.pi)
+    ax.set_theta_zero_location('N')
+    plt.colorbar(im, ax=ax, pad=0.1, shrink=0.8, label='Energy gradient [ergs/cm$^3$]')
+    plt.savefig(os.path.join(results_path, f'energy_gradient_polar_{os.path.basename(file)}.jpg'), dpi=300)
+    plt.close()
diff --git a/nf2/evaluation/spherical/free_energy.py b/nf2/evaluation/spherical/free_energy.py
new file mode 100644
index 0000000..68b362f
--- /dev/null
+++ b/nf2/evaluation/spherical/free_energy.py
@@ -0,0 +1,206 @@
+import os
+
+import numpy as np
+import pfsspy
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, SymLogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map, all_coordinates_from_map
+
+from nf2.data.util import cartesian_to_spherical, vector_cartesian_to_spherical
+from nf2.evaluation.output import SphericalOutput
+
+synoptic_file = '/glade/work/rjarolim/data/global/fd_2283/hmi.synoptic_mr_polfil_720s.2282.Mr_polfil.fits'
+full_disc_file = '/glade/work/rjarolim/data/global/fd_2283/full_disk/hmi.b_720s.20240423_021200_TAI.Br.fits'
+nf2_file = '/glade/work/rjarolim/nf2/spherical/2283_v01/extrapolation_result.nf2'
+
+results_path = '/glade/work/rjarolim/nf2/spherical/2283_v01/results'
+os.makedirs(results_path, exist_ok=True)
+
+synoptic_map = Map(synoptic_file).resample([360, 180] * u.pix)
+full_disc_map = Map(full_disc_file)
+
+potential_r_map = full_disc_map.reproject_to(synoptic_map.wcs)
+mask = np.isnan(potential_r_map.data)
+potential_r_map.data[mask] = synoptic_map.data[mask]
+
+model = SphericalOutput(nf2_file)
+
+latitude_range = 0.398 - np.pi/2, 2.911 - np.pi/2
+longitude_range = 1.673, 4.501
+
+coords = np.stack(np.meshgrid(
+np.linspace(1, 1.3, 128),
+    np.linspace(*latitude_range, np.rad2deg(latitude_range[1] - latitude_range[0]).astype(int) * 2),
+    np.linspace(*longitude_range, np.rad2deg(longitude_range[1] - longitude_range[0]).astype(int) * 2), indexing='ij'),
+    -1)
+
+# PFSS extrapolation
+spherical_boundary_coords = SkyCoord(lon=coords[..., 2] * u.rad, lat=(coords[..., 1]) * u.rad,
+                                     radius=coords[..., 0] * u.solRad, frame=potential_r_map.coordinate_frame)
+
+model_out = model.load_spherical_coords(spherical_boundary_coords)
+b = model_out['b'].to_value(u.G)
+b = vector_cartesian_to_spherical(b, model_out['spherical_coords'])
+j = model_out['j'].to_value(u.G / u.Mm)
+j = vector_cartesian_to_spherical(j, model_out['spherical_coords'])
+
+# potential_r_map.data[np.isnan(potential_r_map.data)] = 0
+
+pfss_in = pfsspy.Input(potential_r_map, 128, 2.5)
+pfss_out = pfsspy.pfss(pfss_in)
+
+potential_shape = spherical_boundary_coords.shape  # required workaround for pfsspy spherical reshape
+spherical_boundary_values = pfss_out.get_bvec(spherical_boundary_coords.reshape((-1,)))
+spherical_boundary_values = spherical_boundary_values.reshape((*potential_shape, 3)).value
+spherical_boundary_values[..., 1] *= -1  # flip B_theta
+potential_b = np.stack([spherical_boundary_values[..., 0],
+                        spherical_boundary_values[..., 1],
+                        spherical_boundary_values[..., 2]], -1)
+
+# compute free energy
+free_energy = (b ** 2).sum(-1) - (potential_b ** 2).sum(-1)
+
+extent = None#[0, 360, -90, 90]
+# plot free energy
+fig, ax = plt.subplots(figsize=(10, 5))
+im = ax.imshow(free_energy.sum(0), origin='lower',
+               norm=LogNorm(vmin=1, vmax=1e7),
+               cmap='cividis', extent=extent)
+ax.set_title('Free Energy')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+# grid with steps of 20
+ax.grid()
+# ax.set_xticks(np.arange(0, 360, 20))
+# ax.set_yticks(np.arange(-80, 90, 20))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+plt.savefig(os.path.join(results_path,
+                         f'free_energy.jpg'),
+            dpi=300)
+plt.close(fig)
+
+# free energy v2
+free_energy = np.linalg.norm(b - potential_b, axis=-1)
+
+extent = None#[0, 360, -90, 90]
+# plot free energy
+fig, ax = plt.subplots(figsize=(10, 5))
+im = ax.imshow(free_energy.sum(0), origin='lower',
+               norm=LogNorm(vmin=10, vmax=1e4),
+               cmap='cividis', extent=extent)
+ax.set_title('Free Energy')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+# grid with steps of 20
+ax.grid()
+# ax.set_xticks(np.arange(0, 360, 20))
+# ax.set_yticks(np.arange(-80, 80, 20))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+plt.savefig(os.path.join(results_path, f'free_energy_v2.jpg'),
+            dpi=300)
+plt.close(fig)
+
+# plot integrated currents
+fig, axs = plt.subplots(1, 4, figsize=(15, 5)) # horizontal
+
+
+ax = axs[0]
+im = ax.imshow(b[0, :, :, 0], origin='lower', vmin=-500, vmax=500, cmap='gray', extent=extent)
+ax.set_title('Photospheric Magnetic Field')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[1]
+current_density = (j[..., :] ** 2).sum(-1) ** 0.5
+j_norm = LogNorm(vmin=1, vmax=1e3)
+im = ax.imshow(current_density.sum(0), origin='lower',
+               norm=j_norm,
+               cmap='plasma', extent=extent)
+ax.set_title('Current Density')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[2]
+current_density = (j[..., 1:] ** 2).sum(-1) ** 0.5
+im = ax.imshow(current_density.sum(0), origin='lower',
+               norm=j_norm,
+               cmap='plasma', extent=extent)
+ax.set_title('Horizontal Current Density')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+
+ax = axs[3]
+current_density = (j[..., 0] ** 2) ** 0.5
+im = ax.imshow(current_density.sum(0), origin='lower',
+               norm=j_norm,
+               cmap='plasma', extent=extent)
+ax.set_title('Radial Current Density')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+
+plt.savefig(os.path.join(results_path, f'currents.jpg'),
+            dpi=300)
+plt.close(fig)
+
+
+
+
+
+# plot photospheric magnetic field
+fig, ax = plt.subplots(figsize=(10, 5))
+im = ax.imshow(b[0, :, :, 0], origin='lower',
+               vmin=-500, vmax=500,
+               cmap='gray', extent=extent)
+ax.set_title('Photospheric Magnetic Field')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+# ax.set_xticks(np.arange(0, 360, 20))
+# ax.set_yticks(np.arange(-80, 80, 20))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+plt.savefig(os.path.join(results_path, f'b.jpg'),
+            dpi=300)
+plt.close(fig)
+
+# plot integrated magnetic field
+fig, ax = plt.subplots(figsize=(10, 5))
+energy = (b ** 2).sum(-1) ** 0.5
+im = ax.imshow(energy.sum(0), origin='lower',
+               norm=LogNorm(vmin=1e1, vmax=1e4),
+               cmap='viridis', extent=extent)
+ax.set_title('Integrated Magnetic Field')
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+ax.grid()
+# ax.set_xticks(np.arange(0, 360, 20))
+# ax.set_yticks(np.arange(-80, 80, 20))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical')
+plt.savefig(os.path.join(results_path, f'energy.jpg'),
+            dpi=300)
+plt.close(fig)
diff --git a/nf2/evaluation/spherical/open_fields.py b/nf2/evaluation/spherical/open_fields.py
new file mode 100644
index 0000000..586f52e
--- /dev/null
+++ b/nf2/evaluation/spherical/open_fields.py
@@ -0,0 +1,123 @@
+import argparse
+import os
+from urllib.request import urlretrieve
+
+import astropy.units as u
+import matplotlib.pyplot as plt
+import numpy as np
+import pfsspy
+import torch
+from astropy import constants
+from astropy.coordinates import SkyCoord
+from astropy.nddata import block_reduce
+from pfsspy import tracing
+from sunpy.map import Map, all_coordinates_from_map, make_heliographic_header
+from sunpy.visualization.colormaps import cm
+
+from nf2.evaluation.output import SphericalOutput
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+parser.add_argument('--out_path', type=str, help='path to the target VTK files')
+parser.add_argument('--cr', type=int, help='Carrington rotation number')
+
+args = parser.parse_args()
+nf2_path = args.nf2_path
+out_path = args.out_path
+
+os.makedirs(out_path, exist_ok=True)
+
+# download AIA synoptic map from NJIT server
+aia_synoptic_path = os.path.join(out_path, 'aia_synoptic.fits')
+if not os.path.exists(aia_synoptic_path):
+    urlretrieve(f'https://sun.njit.edu/coronal_holes/data/aia193_synmap_cr{args.cr}.fits', aia_synoptic_path)
+
+
+def _field_line_endpoint(field_lines):
+    field_lines = np.array([np.abs(fline._r[0] - fline._r[-1]) if len(fline._r) > 0 else 0 for fline in field_lines])
+    return field_lines + 1
+
+
+
+mag_r_map = Map("/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits")#f'/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.{args.cr}.Mr_polfil.fits')
+mag_r_map = mag_r_map.resample([180, 90] * u.pix)
+nrho = 100
+rss = 1.2
+pfss_in = pfsspy.Input(mag_r_map, nrho, rss)
+pfss_out = pfsspy.pfss(pfss_in)
+
+r = constants.R_sun
+coords = all_coordinates_from_map(mag_r_map)
+seeds = coords.ravel()
+seeds = SkyCoord(lat=seeds.lat, lon=seeds.lon, radius=r, frame=coords.frame)
+
+# NF2 field solution
+print('Tracing field lines from NF2...')
+nf2_out = SphericalOutput(nf2_path)
+field_lines = nf2_out.trace(seeds)
+field_map = np.array([np.abs(np.linalg.norm(fline[0]) - np.linalg.norm(fline[-1])) for fline in field_lines]).reshape(coords.shape)
+field_map += 1
+
+# potential field solution
+print('Tracing field lines from PFSS...')
+tracer = tracing.FortranTracer(max_steps=5000)
+field_lines = tracer.trace(seeds, pfss_out)
+# potential_field_map = field_lines.polarities.reshape(coords.shape)
+potential_field_map = _field_line_endpoint(field_lines).reshape(coords.shape)
+
+
+
+# load AIA synoptic map
+aia_map = Map(aia_synoptic_path)
+shape = aia_map.data.shape
+carr_header = make_heliographic_header(aia_map.date, aia_map.observer_coordinate, shape,
+                                       frame='carrington', map_center_longitude=180 * u.deg)
+aia_map = Map(aia_map.data, carr_header)
+aia_map = aia_map.reproject_to(mag_r_map.wcs)
+
+# plot field map with projection
+fig, axs = plt.subplots(5, 1, figsize=(12, 24), subplot_kw={'projection': aia_map})
+
+im = axs[0].imshow(aia_map.data, origin='lower', cmap=cm.sdoaia193)
+plt.colorbar(im, label='Intensity [DN/s]')
+
+cmap = plt.get_cmap('gray')
+cmap.set_bad('yellow')
+im = axs[1].imshow(aia_map.data, origin='lower', cmap=cm.sdoaia193)
+plt.colorbar(im, label='Intensity [DN/s]')
+axs[1].contourf(field_map, levels=[1.05, 1.5], colors=['blue'], origin='lower', alpha=0.3)
+axs[1].contourf(potential_field_map, levels=[1.05, 1.5], colors=['red'], origin='lower', alpha=0.3)
+
+im = axs[2].imshow(field_map, cmap=cmap, origin='lower', vmin=1, vmax=rss)
+plt.colorbar(im, label='Endpoint [$R_{\odot}$]')
+
+im = axs[3].imshow(potential_field_map, cmap=cmap, origin='lower', vmin=1, vmax=rss)
+plt.colorbar(im, label='Endpoint [$R_{\odot}$]')
+
+im = axs[4].imshow(mag_r_map.data, origin='lower', cmap=cmap, vmin=-500, vmax=500)
+plt.colorbar(im, label='Magnetic Field Strength [G]')
+
+# outline bright and dark regions
+# smoothed_aia = block_reduce(aia_map.data, (4, 4), np.nanmean)
+# levels = [np.nanpercentile(smoothed_aia, 10), np.nanpercentile(smoothed_aia, 80)]
+# axs[0].contour(smoothed_aia, levels=levels,
+#                colors=['blue', 'red'], extent=extent, linewidths=3, origin='lower')
+# axs[1].contour(smoothed_aia, levels=levels,
+#                colors=['blue', 'red'], extent=extent, linewidths=3, origin='lower')
+# axs[2].contour(smoothed_aia, levels=levels,
+#                colors=['blue', 'red'], extent=extent, linewidths=3, origin='lower')
+# axs[3].contour(smoothed_aia, levels=levels,
+#                colors=['blue', 'red'], extent=extent, linewidths=3, origin='lower')
+
+# axis labels
+[ax.set_xlabel('Carrington Longitude [deg]') for ax in axs]
+[ax.set_ylabel('Carrington Latitude [deg]') for ax in axs]
+axs[0].set_title('AIA 193')
+axs[1].set_title('Overlay')
+axs[2].set_title('NF2')
+axs[3].set_title('PFSS')
+axs[4].set_title('HMI')
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, f'open_fields.jpg'))
+plt.close()
diff --git a/nf2/evaluation/spherical/potential_vtk.py b/nf2/evaluation/spherical/potential_vtk.py
new file mode 100644
index 0000000..2a79e9b
--- /dev/null
+++ b/nf2/evaluation/spherical/potential_vtk.py
@@ -0,0 +1,116 @@
+import argparse
+import asyncio
+import glob
+import os
+import threading
+
+import numpy as np
+import pfsspy
+import torch
+from astropy.coordinates import SkyCoord
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from torch import nn
+
+from nf2.data.util import vector_cartesian_to_spherical, cartesian_to_spherical, spherical_to_cartesian, \
+    vector_spherical_to_cartesian
+from nf2.evaluation.unpack import load_coords
+from nf2.evaluation.vtk import save_vtk
+from astropy import units as u
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--Br', type=str, help='path to the source radial magnetic field file')
+parser.add_argument('--out_path', type=str, help='path to the target VTK files')
+
+parser.add_argument('--radius_range', nargs='+', type=float, default=(0.999, 2.0), required=False)
+parser.add_argument('--latitude_range', nargs='+', type=float, default=(0 * np.pi, 1 * np.pi), required=False)
+parser.add_argument('--longitude_range', nargs='+', type=float, default=(0 * np.pi, 2 * np.pi), required=False)
+parser.add_argument('--pixels_per_solRad', type=int, default=32, required=False)
+
+args = parser.parse_args()
+out_path = args.out_path
+
+os.makedirs(out_path, exist_ok=True)
+
+radius_range = tuple(args.radius_range)
+latitude_range = tuple(args.latitude_range)
+longitude_range = tuple(args.longitude_range)
+pixels_per_solRad = args.pixels_per_solRad
+
+assert len(radius_range) == 2, 'radius_range must be a tuple of length 2'
+assert len(latitude_range) == 2, 'latitude_range must be a tuple of length 2'
+assert len(longitude_range) == 2, 'longitude_range must be a tuple of length 2'
+
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+mag_r_map = Map(args.Br)
+mag_r_map = mag_r_map.resample([360 * 4, 180 * 4] * u.pix)
+nrho = 9
+rss = 2.5
+pfss_in = pfsspy.Input(mag_r_map, nrho, rss)
+pfss_out = pfsspy.pfss(pfss_in)
+
+fig, axs = plt.subplots(10, 3, figsize=(10, 10))
+b_potential = pfss_out.bg
+for i in range(10):
+    v_min_max = np.abs(b_potential[:, :, i]).max()
+    v_min_max = 500 if v_min_max > 500 else v_min_max
+    axs[i, 0].imshow(b_potential[:, :, i, 0].T, vmin=-v_min_max, vmax=v_min_max, cmap='gray', origin='lower')
+    axs[i, 1].imshow(b_potential[:, :, i, 1].T, vmin=-v_min_max, vmax=v_min_max, cmap='gray', origin='lower')
+    im = axs[i, 2].imshow(b_potential[:, :, i, 2].T, vmin=-v_min_max, vmax=v_min_max, cmap='gray', origin='lower')
+    divider = make_axes_locatable(axs[i, 2])
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    plt.colorbar(im, cax=cax)
+
+[ax.axis('off') for ax in axs.ravel()]
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'potential.jpg'), dpi=300)
+plt.close()
+
+
+vtk_path = os.path.join(out_path, 'potential.vtk')
+
+
+
+spherical_bounds = np.stack(
+    np.meshgrid(np.linspace(radius_range[0], radius_range[1], 50),
+                np.linspace(latitude_range[0], latitude_range[1], 50),
+                np.linspace(longitude_range[0], longitude_range[1], 50), indexing='ij'), -1)
+
+cartesian_bounds = spherical_to_cartesian(spherical_bounds)
+
+x_min, x_max = cartesian_bounds[..., 0].min(), cartesian_bounds[..., 0].max()
+y_min, y_max = cartesian_bounds[..., 1].min(), cartesian_bounds[..., 1].max()
+z_min, z_max = cartesian_bounds[..., 2].min(), cartesian_bounds[..., 2].max()
+
+x_min, x_max = -2, 2
+y_min, y_max = -2, 2
+z_min, z_max = -2, 2
+
+coords = np.stack(
+    np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixels_per_solRad)),
+                np.linspace(y_min, y_max, int((y_max - y_min) * pixels_per_solRad)),
+                np.linspace(z_max, z_min, int((z_max - z_min) * pixels_per_solRad)), indexing='ij'), -1)
+# flipped z axis
+radius = np.sqrt(np.sum(coords ** 2, -1))
+
+spherical_coords = cartesian_to_spherical(coords)
+condition = (spherical_coords[..., 0] >= radius_range[0]) & (spherical_coords[..., 0] < radius_range[1]) \
+            & (spherical_coords[..., 1] > latitude_range[0]) & (spherical_coords[..., 1] < latitude_range[1]) \
+    # & (spherical_coords[..., 2] > longitude_range[0]) & (spherical_coords[..., 2] < longitude_range[1])
+spherical_coords[..., 1] -= np.pi / 2
+sub_coords = spherical_coords[condition]
+
+
+cube_shape = coords.shape[:-1]
+sky_sub_coords = SkyCoord(lon=sub_coords[..., 2] * u.rad,
+                      lat=(sub_coords[..., 1]) * u.rad,
+                      radius=sub_coords[..., 0] * u.solRad, frame=mag_r_map.coordinate_frame)
+sub_b = pfss_out.get_bvec(sky_sub_coords, out_type='cartesian')
+# sub_b[..., 2] *= -1  # flip z axis
+
+b = np.zeros(cube_shape + (3,))
+b[condition] = sub_b
+
+save_vtk(vtk_path, {'B': b}, {'radius': radius})
diff --git a/nf2/evaluation/spherical/traces.py b/nf2/evaluation/spherical/traces.py
new file mode 100644
index 0000000..6e56d53
--- /dev/null
+++ b/nf2/evaluation/spherical/traces.py
@@ -0,0 +1,104 @@
+import argparse
+import os
+
+import astropy.units as u
+import matplotlib.pyplot as plt
+import numpy as np
+import pfsspy
+from astropy import constants
+from astropy.coordinates import SkyCoord
+from pfsspy import tracing
+from sunpy.map import Map, all_coordinates_from_map
+
+from nf2.evaluation.global_field.common import NF2Output
+
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--nf2_path', type=str, help='path to the source NF2 file')
+parser.add_argument('--out_path', type=str, help='path to the target VTK files')
+parser.add_argument('--cr', type=int, help='Carrington rotation number')
+
+args = parser.parse_args()
+nf2_path = args.nf2_path
+out_path = args.out_path
+max_radius = 1.3
+
+os.makedirs(out_path, exist_ok=True)
+
+
+def _field_line_endpoint(field_lines):
+    field_lines = np.array([np.abs(fline._r[0] - fline._r[-1]) if len(fline._r) > 0 else 0 for fline in field_lines])
+    return field_lines + 1
+
+
+# potential field solution
+mag_r_map = Map(f'/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.{args.cr}.Mr_polfil.fits')
+mag_r_map = mag_r_map.resample([360, 180] * u.pix)
+nrho = 100
+rss = 2.5
+pfss_in = pfsspy.Input(mag_r_map, nrho, rss)
+pfss_out = pfsspy.pfss(pfss_in)
+
+r = constants.R_sun
+
+# select equally distributed seeds
+coords = all_coordinates_from_map(mag_r_map).ravel()
+values = np.abs(mag_r_map.data).ravel()
+sort_idx = np.argsort(values)
+X2 = values[sort_idx]
+
+full_idx = np.linspace(0, 1, len(X2))
+selected_idx = np.linspace(0, 1, 4096)
+selected_idx = np.argmin(np.abs(full_idx[None, :] - selected_idx[:, None]), axis=1)
+
+seeds = coords[sort_idx][selected_idx]
+seeds = SkyCoord(lat=seeds.lat, lon=seeds.lon, radius=r, frame=seeds.frame)
+
+tracer = tracing.FortranTracer(max_steps=1000)
+potential_field_lines = tracer.trace(seeds, pfss_out)
+
+nf2_out = NF2Output(nf2_path, pfss_in)
+field_lines = tracer.trace(seeds, nf2_out)
+
+# plot 3D field lines
+fig = plt.figure(figsize=(20, 10))
+
+
+# u, v = np.mgrid[0:2 * np.pi:radial_b_map.shape[1] * 1j, 0:np.pi:radial_b_map.shape[0] * 1j]
+# x = 1 * np.cos(u) * np.sin(v)
+# y = 1 * np.sin(u) * np.sin(v)
+# z = 1 * np.cos(v)
+# color_mapping = plt.get_cmap('gray')(Normalize(-500, 500)(radial_b_map))
+# ax.plot_surface(x.T, y.T, z.T, facecolors=color_mapping, cstride=2, rstride=2, zorder=10)
+
+ax = fig.add_subplot(121, projection='3d')
+for field_line in potential_field_lines:
+    color = {0: 'black', -1: 'tab:blue', 1: 'tab:red'}.get(field_line.polarity)
+    coords = field_line.coords
+    coords.representation_type = 'cartesian'
+    radius = np.sqrt(coords.x ** 2 + coords.y ** 2 + coords.z ** 2) / constants.R_sun
+    coords.x[radius > 1.3] = np.nan
+    coords.x[coords.x > 0] = np.nan
+    ax.plot(coords.x / constants.R_sun,
+            coords.y / constants.R_sun,
+            coords.z / constants.R_sun,
+            color=color, linewidth=1)
+
+ax.view_init(azim=0, elev=0)
+
+ax = fig.add_subplot(122, projection='3d')
+for field_line in field_lines:
+    color = {0: 'black', -1: 'tab:blue', 1: 'tab:red'}.get(field_line.polarity)
+    coords = field_line.coords
+    coords.representation_type = 'cartesian'
+    radius = np.sqrt(coords.x ** 2 + coords.y ** 2 + coords.z ** 2) / constants.R_sun
+    coords.x[radius > 1.3] = np.nan
+    coords.x[coords.x > 0] = np.nan
+    ax.plot(coords.x / constants.R_sun,
+            coords.y / constants.R_sun,
+            coords.z / constants.R_sun,
+            color=color, linewidth=1)
+
+ax.view_init(azim=0, elev=0)
+
+plt.savefig(os.path.join(out_path, f'field_lines.jpg'), dpi=300)
+plt.close()
diff --git a/nf2/evaluation/sst/__init__.py b/nf2/evaluation/sst/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/evaluation/sst/compare_disambiguation.py b/nf2/evaluation/sst/compare_disambiguation.py
new file mode 100644
index 0000000..8138dc3
--- /dev/null
+++ b/nf2/evaluation/sst/compare_disambiguation.py
@@ -0,0 +1,75 @@
+import argparse
+import os
+
+import numpy as np
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+
+from nf2.data.util import los_trv_azi_to_img, img_to_los_trv_azi
+from nf2.evaluation.output import HeightTransformOutput
+
+parser = argparse.ArgumentParser(description='Evaluate SHARP snapshot.')
+parser.add_argument('--nf2_path', type=str, help='path to the nf2 project.')
+parser.add_argument('--b_los', type=str, help='path to the FITS file containing the line-of-sight magnetic field.')
+parser.add_argument('--b_trv', type=str, help='path to the FITS file containing the transverse magnetic field.')
+parser.add_argument('--b_azi', type=str, help='path to the FITS file containing the azimuthal magnetic field.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+nf2_out = HeightTransformOutput(args.nf2_path)
+
+b_nf2 = nf2_out.load_cube(height_range=(0, 1))['B']
+
+b_los = Map(args.b_los).data
+b_trv = Map(args.b_trv).data
+b_azi = Map(args.b_azi).data
+
+b = np.stack([b_los, b_trv, np.pi - b_azi], axis=-1)
+b_xyz = los_trv_azi_to_img(b)
+
+print(np.nansum(b - img_to_los_trv_azi(b_xyz)))
+
+fig, axs = plt.subplots(2, 3, figsize=(16, 4))
+
+ax = axs[0, 0]
+im = ax.imshow(b_nf2[:, :, 0, 0].value.T, vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='Bx [G]')
+ax.set_title('NF2')
+#
+ax = axs[0, 1]
+im = ax.imshow(b_nf2[:, :, 0, 1].value.T, vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='By [G]')
+
+ax = axs[0, 2]
+im = ax.imshow(b_nf2[:, :, 0, 2].value.T, vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='Bz [G]')
+
+ax = axs[1, 0]
+im = ax.imshow(b_xyz[:, :, 0], vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='Bx [G]')
+ax.set_title('SHARP')
+
+ax = axs[1, 1]
+im = ax.imshow(b_xyz[:, :, 1], vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='By [G]')
+
+ax = axs[1, 2]
+im = ax.imshow(b_xyz[:, :, 2], vmin=-1000, vmax=1000, cmap='gray', origin='lower')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, label='Bz [G]')
+
+fig.tight_layout()
+fig.savefig(os.path.join(args.output, f'compare_disambiguation.jpg'))
+plt.close(fig)
diff --git a/nf2/evaluation/unpack.py b/nf2/evaluation/unpack.py
index b9e1119..8b8d70a 100644
--- a/nf2/evaluation/unpack.py
+++ b/nf2/evaluation/unpack.py
@@ -8,7 +8,7 @@
 from torch import nn
 from tqdm import tqdm
 
-from nf2.train.model import VectorPotentialModel
+from nf2.train.model import VectorPotentialModel, calculate_current
 
 
 def load_cube(save_path, device=None, z=None, strides=1, **kwargs):
@@ -19,9 +19,10 @@ def load_cube(save_path, device=None, z=None, strides=1, **kwargs):
     cube_shape = state['cube_shape']
     z = z if z is not None else cube_shape[2]
     coords = np.stack(np.mgrid[:cube_shape[0]:strides, :cube_shape[1]:strides, :z:strides], -1)
-    return load_coords(model, cube_shape, state['spatial_norm'],
+    return load_coords(model, state['spatial_norm'],
                        state['b_norm'], coords, device, **kwargs)
 
+
 def load_height_surface(save_path, device=None, strides=1, batch_size=1000, progress=False, **kwargs):
     if device is None:
         device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -62,6 +63,7 @@ def load_height_surface(save_path, device=None, strides=1, batch_size=1000, prog
     slices = slices.view(*coords_shape).numpy()
     return slices
 
+
 def load_height_cube(save_path, *args, device=None, strides=1, **kwargs):
     if device is None:
         device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -73,9 +75,10 @@ def load_height_cube(save_path, *args, device=None, strides=1, **kwargs):
     cube = np.stack(np.mgrid[:cube_shape[0]:strides, :cube_shape[1]:strides, :cube_shape[2]:strides], -1)[..., 2]
     contour_cube = np.zeros_like(cube)
     for i in range(slice.shape[2]):
-        contour_cube[cube > slice[:, :, i:i+1, 2]] = i + 1
+        contour_cube[cube > slice[:, :, i:i + 1, 2]] = i + 1
     return contour_cube
 
+
 def load_shape(save_path, device=None):
     if device is None:
         device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -91,7 +94,7 @@ def load_slice(save_path, z=0, device=None, **kwargs):
     cube_shape = state['cube_shape']
     z = z if z is not None else cube_shape[2]
     coords = np.stack(np.mgrid[:cube_shape[0], :cube_shape[1], z:z + 1], -1)
-    return load_coords(model, cube_shape, state['spatial_norm'],
+    return load_coords(model, state['spatial_norm'],
                        state['b_norm'], coords, device, **kwargs)
 
 
@@ -101,41 +104,55 @@ def load_coords_from_state(save_path, coords, device=None, **kwargs):
     state = torch.load(save_path, map_location=device)
     model = nn.DataParallel(state['model'])
     cube_shape = state['cube_shape']
-    return load_coords(model, cube_shape, state['spatial_norm'], state['b_norm'], coords, device,
-                       **kwargs)
-
-
-def load_coords(model, cube_shape, spatial_norm, b_norm, coords, device, batch_size=1000, progress=False):
     assert np.all(coords[..., 0] < cube_shape[0]), 'Invalid x coordinate, maximum is %d' % cube_shape[0]
     assert np.all(coords[..., 1] < cube_shape[1]), 'Invalid x coordinate, maximum is %d' % cube_shape[1]
     assert np.all(coords[..., 2] < cube_shape[2]), 'Invalid x coordinate, maximum is %d' % cube_shape[2]
+    return load_coords(model, state['spatial_norm'], state['b_norm'], coords, device,
+                       **kwargs)
+
 
+def load_coords(model, spatial_norm, b_norm, coords, device, batch_size=1000, progress=False, compute_currents=False):
     def _load(coords):
         # normalize and to tensor
         coords = torch.tensor(coords / spatial_norm, dtype=torch.float32)
         coords_shape = coords.shape
-        coords = coords.view((-1, 3))
+        coords = coords.reshape((-1, 3))
 
         cube = []
+        j_cube = []
         it = range(int(np.ceil(coords.shape[0] / batch_size)))
         it = tqdm(it) if progress else it
         for k in it:
+            model.zero_grad()
             coord = coords[k * batch_size: (k + 1) * batch_size]
             coord = coord.to(device)
             coord.requires_grad = True
-            cube += [model(coord).detach().cpu()]
+            result = model(coord)
+            b_batch = result['b']
+            if compute_currents:
+                j_batch = calculate_current(b_batch, coord)
+                j_cube += [j_batch.detach().cpu()]
+            cube += [b_batch.detach().cpu()]
 
         cube = torch.cat(cube)
-        cube = cube.view(*coords_shape).numpy()
+        cube = cube.reshape(*coords_shape).numpy()
         b = cube * b_norm
+        if compute_currents:
+            j_cube = torch.cat(j_cube)
+            j_cube = j_cube.reshape(*coords_shape).numpy()
+            j = j_cube * b_norm / spatial_norm
+            return b, j
         return b
-    if isinstance(model, VectorPotentialModel) or \
-            (isinstance(model, nn.DataParallel) and isinstance(model.module, VectorPotentialModel)):
+
+    if (compute_currents or
+            isinstance(model, VectorPotentialModel) or (
+                    isinstance(model, nn.DataParallel) and isinstance(model.module, VectorPotentialModel))):
         return _load(coords)
     else:
         with torch.no_grad():
             return _load(coords)
 
+
 def load_B_map(nf2_file, component=2):
     state = torch.load(nf2_file)
     meta = state['meta_data']
diff --git a/nf2/evaluation/vtk.py b/nf2/evaluation/vtk.py
index 79c119e..72dbd0c 100644
--- a/nf2/evaluation/vtk.py
+++ b/nf2/evaluation/vtk.py
@@ -2,33 +2,50 @@
 from tvtk.api import tvtk, write_data
 
 
-def save_vtk(vec, path, name, scalar=None, scalar_name='scalar', Mm_per_pix=720e-3):
+def save_vtk(path, coords=None, vectors={}, scalars={}, Mm_per_pix=720e-3):
     """Save numpy array as VTK file
 
-    :param vec: numpy array of the vector field (x, y, z, c)
+    :param vectors: numpy array of the vector field (x, y, z, c)
     :param path: path to the target VTK file
     :param name: label of the vector field (e.g., B)
     :param Mm_per_pix: pixel size in Mm. 360e-3 for original HMI resolution. (default bin2 pixel scale)
     """
     # Unpack
-    dim = vec.shape[:-1]
-    # Generate the grid
-    pts = np.stack(np.mgrid[0:dim[0], 0:dim[1], 0:dim[2]], -1).astype(np.int64) * Mm_per_pix
-    # reorder the points and vectors in agreement with VTK
-    # requirement of x first, y next and z last.
-    pts = pts.transpose(2, 1, 0, 3)
-    pts = pts.reshape((-1, 3))
-    vectors = vec.transpose(2, 1, 0, 3)
-    vectors = vectors.reshape((-1, 3))
+    if len(vectors) > 0:
+        dim = list(vectors.values())[0].shape[:-1]
+    elif len(scalars) > 0:
+        dim = list(scalars.values())[0].shape
+    else:
+        raise ValueError('No data to save')
+
+    if coords is None:
+        # Generate the grid
+        pts = np.stack(np.mgrid[0:dim[0], 0:dim[1], 0:dim[2]], -1).astype(np.int64) * Mm_per_pix
+        # reorder the points and vectors in agreement with VTK
+        pts = pts.transpose(2, 1, 0, 3)
+        pts = pts.reshape((-1, 3))
+    else:
+        pts = coords
+        # reorder the points and vectors in agreement with VTK
+        pts = pts.transpose(2, 1, 0, 3)
+        pts = pts.reshape((-1, 3))
+
 
     sg = tvtk.StructuredGrid(dimensions=dim, points=pts)
-    sg.point_data.vectors = vectors
-    sg.point_data.vectors.name = name
-    if scalar is not None:
-        scalars = scalar.transpose(2, 1, 0)
-        scalars = scalars.reshape((-1))
-        sg.point_data.add_array(scalars)
-        sg.point_data.get_array(1).name = scalar_name
+    i = 0
+    for v_name, v in vectors.items():
+        v = v.transpose(2, 1, 0, 3)
+        v = v.reshape((-1, 3))
+        sg.point_data.add_array(v)
+        sg.point_data.get_array(i).name = v_name
+        sg.point_data.update()
+        i += 1
+    for s_name, s in scalars.items():
+        s = s.transpose(2, 1, 0)
+        s = s.reshape((-1))
+        sg.point_data.add_array(s)
+        sg.point_data.get_array(i).name = s_name
         sg.point_data.update()
+        i += 1
 
     write_data(sg, path)
diff --git a/nf2/extrapolate.py b/nf2/extrapolate.py
index 0db769f..81db90c 100644
--- a/nf2/extrapolate.py
+++ b/nf2/extrapolate.py
@@ -1,90 +1,129 @@
 import argparse
-import json
 import os
 import shutil
 
 import torch
+import yaml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint, LambdaCallback
 from pytorch_lightning.loggers import WandbLogger
 
+from nf2.loader.analytical import AnalyticDataModule
+from nf2.loader.fits import FITSDataModule
+from nf2.loader.general import NumpyDataModule
+from nf2.loader.muram import MURaMDataModule
+from nf2.loader.spherical import SphericalDataModule
+from nf2.loader.vsm import VSMDataModule
+from nf2.train.mapping import load_callbacks
 from nf2.train.module import NF2Module, save
-from nf2.train.data_loader import NumpyDataModule, SOLISDataModule, FITSDataModule, AnalyticDataModule, SHARPDataModule
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--config', type=str, required=True,
-                    help='config file for the simulation')
-args = parser.parse_args()
-
-with open(args.config) as config:
-    info = json.load(config)
-    for key, value in info.items():
-        args.__dict__[key] = value
-
-base_path = args.base_path
-os.makedirs(base_path, exist_ok=True)
-
-save_path = os.path.join(base_path, 'extrapolation_result.nf2')
-
-# init logging
-wandb_id = args.logging['wandb_id'] if 'wandb_id' in args.logging else None
-log_model = args.logging['wandb_log_model'] if 'wandb_log_model' in args.logging else False
-wandb_logger = WandbLogger(project=args.logging['wandb_project'], name=args.logging['wandb_name'], offline=False,
-                           entity=args.logging['wandb_entity'], id=wandb_id, dir=base_path, log_model=log_model)
-wandb_logger.experiment.config.update(vars(args), allow_val_change=True)
-
-# restore model checkpoint from wandb
-if wandb_id is not None:
-    checkpoint_reference = f"{args.logging['wandb_entity']}/{args.logging['wandb_project']}/model-{args.logging['wandb_id']}:latest"
-    artifact = wandb_logger.use_artifact(checkpoint_reference, artifact_type="model")
-    artifact.download(root=base_path)
-    shutil.move(os.path.join(base_path, 'model.ckpt'), os.path.join(base_path, 'last.ckpt'))
-    args.data['plot_overview'] = False  # skip overview plot for restored model
-
-if 'work_directory' not in args.data or args.data['work_directory'] is None:
-    args.data['work_directory'] = base_path
-if args.data["type"] == 'numpy':
-    data_module = NumpyDataModule(**args.data)
-elif args.data["type"] == 'sharp':
-    data_module = SHARPDataModule(**args.data)
-elif args.data["type"] == 'fits':
-    data_module = FITSDataModule(**args.data)
-elif args.data["type"] == 'solis':
-    data_module = SOLISDataModule(**args.data)
-elif args.data["type"] == 'analytical':
-    data_module = AnalyticDataModule(**args.data)
-else:
-    raise NotImplementedError(f'Unknown data loader {args.data["type"]}')
-
-validation_settings = {'cube_shape': data_module.cube_dataset.coords_shape,
-                       'gauss_per_dB': args.data["b_norm"],
-                       'Mm_per_ds': args.data["Mm_per_pixel"] * args.data["spatial_norm"]}
-
-nf2 = NF2Module(validation_settings, **args.model, **args.training)
-
-config = {'data': args.data, 'model': args.model, 'training': args.training}
-save_callback = LambdaCallback(
-    on_validation_end=lambda *args: save(save_path, nf2.model, data_module, config, nf2.height_mapping_model))
-checkpoint_callback = ModelCheckpoint(dirpath=base_path, every_n_train_steps=args.training["validation_interval"],
-                                      save_last=True)
-
-torch.set_float32_matmul_precision('medium')  # for A100 GPUs
-n_gpus = torch.cuda.device_count()
-trainer = Trainer(max_epochs=2,
-                  logger=wandb_logger,
-                  devices=n_gpus if n_gpus > 0 else None,
-                  accelerator='gpu' if n_gpus >= 1 else None,
-                  strategy='dp' if n_gpus > 1 else None,  # ddp breaks memory and wandb
-                  num_sanity_val_steps=0,
-                  val_check_interval=args.training['validation_interval'],
-                  gradient_clip_val=0.1,
-                  callbacks=[checkpoint_callback, save_callback], )
-
-trainer.fit(nf2, data_module, ckpt_path='last')
-save(save_path, nf2.model, data_module, config, height_mapping_model=nf2.height_mapping_model)
-# clean up
-data_module.clear()
-
-
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
+
+
+def run(base_path, data, work_directory=None, logging={}, model={}, training={}, config=None):
+    """Run the simulation with the given configuration.
+
+    This function initializes the data loader, the model, the training loop and the logging.
+    The simulation is run with the given configuration and the results are stored in the base_path.
+    Use the configurations for logging, model and training to overwrite the default settings.
+
+    Args:
+        base_path: Path to the directory where the simulation results are stored.
+        data: Dictionary with the data loader configuration.
+        work_directory: Path to the directory where the data is stored. If None, the base_path is used.
+        logging: Dictionary with the logging configuration.
+        model: Dictionary with the model configuration.
+        training: Dictionary with the training configuration.
+        config: Dictionary with the configuration for the simulation.
+    """
+    os.makedirs(base_path, exist_ok=True)
+
+    if work_directory is None:
+        work_directory = os.path.join(base_path, 'work')
+    os.makedirs(work_directory, exist_ok=True)
+    data['work_directory'] = work_directory # set work directory for data loaders
+
+    save_path = os.path.join(base_path, 'extrapolation_result.nf2')
+
+    # init logging
+    wandb_logger = WandbLogger(**logging, save_dir=work_directory)
+    config_dict = {'base_path': base_path, 'work_directory': work_directory, 'logging': logging,
+                   'model': model, 'training': training, 'config': config}
+    wandb_logger.experiment.config.update(config_dict, allow_val_change=True)
+
+    # restore model checkpoint from wandb
+    if 'id' in logging:
+        checkpoint_reference = f"{logging['entity']}/{logging['project']}/model-{logging['id']}:latest"
+        artifact = wandb_logger.use_artifact(checkpoint_reference, artifact_type="model")
+        artifact.download(root=base_path)
+        shutil.move(os.path.join(base_path, 'model.ckpt'), os.path.join(base_path, 'last.ckpt'))
+        data['plot_overview'] = False  # skip overview plot for restored model
+
+    if data["type"] == 'numpy':
+        data_module = NumpyDataModule(**data)
+    elif data["type"] == 'fits':
+        data_module = FITSDataModule(**data)
+    elif data["type"] == 'solis':
+        data_module = VSMDataModule(**data)
+    elif data["type"] == 'analytical':
+        data_module = AnalyticDataModule(**data)
+    elif data["type"] == 'spherical':
+        data_module = SphericalDataModule(**data)
+    elif data["type"] == 'muram':
+        data_module = MURaMDataModule(**data)
+    else:
+        raise NotImplementedError(f'Unknown data loader {data["type"]}')
+
+    # initialize callbacks
+    callbacks = load_callbacks(data_module)
+
+    nf2 = NF2Module(data_module.validation_dataset_mapping, data_module.config,
+                    model_kwargs=model, **training)
+
+    config_dict = {'data': data, 'model': model, 'training': training, 'config': config}
+    val_check_interval = int(training['validation_interval']) if "validation_interval" in training else None
+    val_every_n_epochs = training['check_val_every_n_epoch'] if 'check_val_every_n_epoch' in training else None
+    max_epochs = int(training['epochs']) if 'epochs' in training else 10
+
+    save_callback = LambdaCallback(
+        on_validation_end=lambda *_: save(save_path, nf2, data_module, config_dict))
+    checkpoint_callback = ModelCheckpoint(dirpath=base_path,
+                                          every_n_train_steps=val_check_interval,
+                                          every_n_epochs=val_every_n_epochs,
+                                          save_last=True)
+
+    torch.set_float32_matmul_precision('medium')  # for A100 GPUs
+    n_gpus = torch.cuda.device_count()
+    callbacks += [checkpoint_callback, save_callback]
+
+    trainer = Trainer(max_epochs=max_epochs,
+                      logger=wandb_logger,
+                      devices=n_gpus if n_gpus > 0 else None,
+                      accelerator='gpu' if n_gpus >= 1 else None,
+                      strategy='dp' if n_gpus > 1 else None,  # ddp breaks memory and wandb
+                      num_sanity_val_steps=0,
+                      val_check_interval=val_check_interval,
+                      check_val_every_n_epoch=val_every_n_epochs,
+                      gradient_clip_val=0.1,
+                      callbacks=callbacks)
+
+    trainer.fit(nf2, data_module, ckpt_path='last')
+    save(save_path, nf2, data_module, config_dict)
+    # clean up
+    data_module.clear()
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True,
+                        help='config file for the simulation')
+    args = parser.parse_args()
+
+    with open(args.config) as config:
+        info = yaml.safe_load(config)
+        for key, value in info.items():
+            args.__dict__[key] = value
+
+    run(**args.__dict__)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/extrapolate_series.py b/nf2/extrapolate_series.py
index 9dda80e..5dc8f6f 100644
--- a/nf2/extrapolate_series.py
+++ b/nf2/extrapolate_series.py
@@ -1,80 +1,172 @@
 import argparse
 import glob
-import json
 import os
+import shutil
 
 import torch
+import yaml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import LambdaCallback
+from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.loggers import WandbLogger
 
+from nf2.loader.fits import FITSSeriesDataModule
+from nf2.loader.spherical import SphericalSeriesDataModule
+from nf2.train.mapping import load_callbacks
 from nf2.train.module import NF2Module, save
-from nf2.train.data_loader import SHARPSeriesDataModule
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--config', type=str, required=True,
-                    help='config file for the simulation')
-args = parser.parse_args()
-
-with open(args.config) as config:
-    info = json.load(config)
-    for key, value in info.items():
-        args.__dict__[key] = value
-
-base_path = args.base_path
-os.makedirs(base_path, exist_ok=True)
-
-# init logging
-wandb_id = args.logging['wandb_id'] if 'wandb_id' in args.logging else None
-wandb_logger = WandbLogger(project=args.logging['wandb_project'], name=args.logging['wandb_name'], offline=False,
-                           entity=args.logging['wandb_entity'], id=wandb_id, dir=base_path, log_model='all')
-wandb_logger.experiment.config.update(vars(args), allow_val_change=True)
-
-
-
-data_path = args.data['paths']
-hmi_p_files = sorted(glob.glob(os.path.join(data_path, '*Bp.fits')))  # x
-hmi_t_files = sorted(glob.glob(os.path.join(data_path, '*Bt.fits')))  # y
-hmi_r_files = sorted(glob.glob(os.path.join(data_path, '*Br.fits')))  # z
-err_p_files = sorted(glob.glob(os.path.join(data_path, '*Bp_err.fits')))  # x
-err_t_files = sorted(glob.glob(os.path.join(data_path, '*Bt_err.fits')))  # y
-err_r_files = sorted(glob.glob(os.path.join(data_path, '*Br_err.fits')))  # z
-
-data_paths = list(zip(hmi_p_files, err_p_files, hmi_t_files, err_t_files, hmi_r_files, err_r_files))
-
-# reload model training
-ckpts = sorted(glob.glob(os.path.join(base_path, '*.nf2')))
-meta_path = ckpts[-1] if len(ckpts) > 0 else args.meta_path  # reload last savepoint
-data_paths = data_paths[len(ckpts):]  # select remaining extrapolations
-
-if 'work_directory' not in args.data or args.data['work_directory'] is None:
-    args.data['work_directory'] = base_path
-data_module = SHARPSeriesDataModule(data_paths, **args.data)
-validation_settings = {'cube_shape': data_module.cube_dataset.coords_shape,
-                       'gauss_per_dB': args.data["b_norm"],
-                       'Mm_per_ds': args.data["Mm_per_pixel"] * args.data["spatial_norm"]}
-nf2 = NF2Module(validation_settings, meta_path=meta_path, **args.model, **args.training)
-
-# callback
-config = {'data': args.data, 'model': args.model, 'training': args.training}
-save_callback = LambdaCallback(on_train_epoch_end=lambda *args: save(
-    os.path.join(base_path, os.path.basename(data_paths[nf2.current_epoch][0]).split('.')[-3] + '.nf2'),
-    nf2.model, data_module, config, nf2.height_mapping_model))
-
-# general training parameters
-torch.set_float32_matmul_precision('medium')  # for A100 GPUs
-n_gpus = torch.cuda.device_count()
-
-trainer = Trainer(max_epochs=len(data_paths),
-                  logger=wandb_logger,
-                  devices=n_gpus,
-                  accelerator='gpu' if n_gpus >= 1 else None,
-                  strategy='dp' if n_gpus > 1 else None,  # ddp breaks memory and wandb
-                  num_sanity_val_steps=0, callbacks=[save_callback],
-                  gradient_clip_val=0.1, reload_dataloaders_every_n_epochs=1,
-                  check_val_every_n_epoch=args.training['check_val_every_n_epoch'] if 'check_val_every_n_epoch' in args.training else 1,)
-trainer.fit(nf2, data_module)
-
-
-def main(): # workaround for entry_points
-    pass
\ No newline at end of file
+
+
+def run(base_path, data, meta_path, work_directory=None, logging={}, model={}, training={}, config=None):
+    """Run the simulation with the given configuration.
+
+    This function initializes the data loader, the model, the training loop and the logging.
+    The simulation is run with the given configuration and the results are stored in the base_path.
+    Use the configurations for logging, model and training to overwrite the default settings.
+
+    Args:
+        base_path: Path to the directory where the simulation results are stored.
+        data: Dictionary with the data loader configuration.
+        meta_path: Path to the initial model checkpoint (run extrapolation first).
+        work_directory: Path to the directory where the data is stored. If None, the base_path is used.
+        logging: Dictionary with the logging configuration.
+        model: Dictionary with the model configuration.
+        training: Dictionary with the training configuration.
+        config: Dictionary with the configuration for the simulation.
+    """
+    os.makedirs(base_path, exist_ok=True)
+
+    if work_directory is None:
+        work_directory = os.path.join(base_path, 'work')
+    os.makedirs(work_directory, exist_ok=True)
+    data['work_directory'] = work_directory  # set work directory for data loaders
+
+    # init logging
+    wandb_logger = WandbLogger(**logging, save_dir=work_directory)
+    config_dict = {'base_path': base_path, 'work_directory': work_directory, 'logging': logging,
+                   'model': model, 'training': training, 'config': config}
+    wandb_logger.experiment.config.update(config_dict, allow_val_change=True)
+
+    # restore model checkpoint from wandb
+    if 'id' in logging:
+        assert 'entity' in logging and 'project' in logging, '"entity" and "project" must be provided to continue from wandb checkpoint'
+        checkpoint_reference = f"{logging['entity']}/{logging['project']}/model-{logging['id']}:latest"
+        artifact = wandb_logger.use_artifact(checkpoint_reference, artifact_type="model")
+        artifact.download(root=base_path)
+        shutil.move(os.path.join(base_path, 'model.ckpt'), os.path.join(base_path, 'last.ckpt'))
+        data['plot_overview'] = False  # skip overview plot for restored model
+
+    fits_paths, error_paths = _load_paths(data)
+
+    # reload model training
+    ckpts = sorted(glob.glob(os.path.join(base_path, '*.nf2')))
+    ckpt_path = 'last' if len(ckpts) > 0 else meta_path  # reload last savepoint
+    fits_paths = fits_paths[len(ckpts):]  # select remaining extrapolations
+    error_paths = error_paths[len(ckpts):] if error_paths is not None else None
+
+    # initialize data module
+    if data["type"] == 'sharp':
+        data_module = FITSSeriesDataModule(fits_paths, error_paths=error_paths, **data)
+    elif data["type"] == 'spherical':
+        data_module = SphericalSeriesDataModule(fits_paths, **data)
+    else:
+        raise NotImplementedError(f'Unknown data loader {data["type"]}')
+
+    callbacks = load_callbacks(data_module)
+
+    nf2 = NF2Module(data_module.validation_dataset_mapping, data_module.config, model_kwargs=model, **training)
+
+    reload_dataloaders_interval = training[
+        'reload_dataloaders_every_n_epochs'] if 'reload_dataloaders_every_n_epochs' in training else 1
+
+    # callback
+    config_dict = {'data': data, 'model': model, 'training': training, 'config': config}
+    save_callback = LambdaCallback(
+        on_train_epoch_end=lambda *args:
+        save(os.path.join(base_path, data_module.current_id + '.nf2'),
+             nf2, data_module, config_dict))
+
+    checkpoint_callback = ModelCheckpoint(dirpath=base_path,
+                                          every_n_epochs=reload_dataloaders_interval,
+                                          save_last=True)
+
+    # general training parameters
+    torch.set_float32_matmul_precision('medium')  # for A100 GPUs
+    n_gpus = torch.cuda.device_count()
+
+    val_check_interval = training['check_val_every_n_epoch'] if 'check_val_every_n_epoch' in training else 1
+    trainer = Trainer(max_epochs=-1,
+                      logger=wandb_logger,
+                      devices=n_gpus,
+                      accelerator='gpu' if n_gpus >= 1 else None,
+                      strategy='dp' if n_gpus > 1 else None,  # ddp breaks memory and wandb
+                      num_sanity_val_steps=0, callbacks=[save_callback, checkpoint_callback, *callbacks],
+                      gradient_clip_val=0.1, reload_dataloaders_every_n_epochs=reload_dataloaders_interval,
+                      check_val_every_n_epoch=val_check_interval)
+    trainer.fit(nf2, data_module, ckpt_path=ckpt_path)
+
+
+def _load_paths(data):
+    data_path = data['data_path']
+    if isinstance(data_path, str):
+        p_files = sorted(glob.glob(os.path.join(data_path, '*Bp.fits')))  # x
+        t_files = sorted(glob.glob(os.path.join(data_path, '*Bt.fits')))  # y
+        r_files = sorted(glob.glob(os.path.join(data_path, '*Br.fits')))  # z
+        err_p_files = sorted(glob.glob(os.path.join(data_path, '*Bp_err.fits')))  # x
+        err_t_files = sorted(glob.glob(os.path.join(data_path, '*Bt_err.fits')))  # y
+        err_r_files = sorted(glob.glob(os.path.join(data_path, '*Br_err.fits')))  # z
+
+        assert len(p_files) == len(t_files) == len(r_files), f'Number of files in data path {data_path} does not match'
+        fits_paths = list(zip(p_files, t_files, r_files))
+        fits_paths = [{'Bp': d[0], 'Bt': d[1], 'Br': d[2]} for d in fits_paths]
+
+        if len(err_p_files) > 0 or len(err_t_files) > 0 or len(err_r_files) > 0:
+            assert len(p_files) == len(err_p_files) == len(t_files) == len(err_t_files) == len(r_files) == len(
+                err_r_files), \
+                f'Number of files in data path {data_path} does not match'
+            error_paths = list(zip(err_p_files, err_t_files, err_r_files))
+            error_paths = [{'Bp_err': d[0], 'Bt_err': d[1], 'Br_err': d[2]} for d in error_paths]
+        else:
+            error_paths = None
+    elif isinstance(data_path, dict):
+        p_files = sorted(glob.glob(data_path['Bp']))  # x
+        t_files = sorted(glob.glob(data_path['Bt']))  # y
+        r_files = sorted(glob.glob(data_path['Br']))  # z
+        err_p_files = sorted(glob.glob(data_path['Bp_err'])) if 'Bp_err' in data_path else None  # x
+        err_t_files = sorted(glob.glob(data_path['Bt_err'])) if 'Bt_err' in data_path else None  # y
+        err_r_files = sorted(glob.glob(data_path['Br_err'])) if 'Br_err' in data_path else None  # z
+
+        if err_p_files is not None and err_t_files is not None and err_r_files is not None:
+            assert len(p_files) == len(err_p_files) == len(t_files) == len(err_t_files) == len(r_files) == len(
+                err_r_files), \
+                f'Number of files in data path {data_path} does not match'
+            fits_paths = list(zip(p_files, t_files, r_files))
+            fits_paths = [{'Bp': d[0], 'Bt': d[1], 'Br': d[2], } for d in fits_paths]
+            error_paths = list(zip(err_p_files, err_t_files, err_r_files))
+            error_paths = [{'Bp_err': d[0], 'Bt_err': d[1], 'Br_err': d[2], } for d in error_paths]
+        else:
+            assert len(p_files) == len(t_files) == len(r_files), \
+                f'Number of files in data path {data_path} does not match'
+            fits_paths = list(zip(p_files, t_files, r_files))
+            fits_paths = [{'Bp': d[0], 'Bt': d[1], 'Br': d[2]} for d in fits_paths]
+            error_paths = None
+    else:
+        raise NotImplementedError(f'Unknown data path type {type(data_path)}')
+    return fits_paths, error_paths
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--config', type=str, required=True,
+                        help='config file for the simulation')
+    args = parser.parse_args()
+
+    with open(args.config) as config:
+        info = yaml.safe_load(config)
+        for key, value in info.items():
+            args.__dict__[key] = value
+
+    run(**args.__dict__)
+
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/nf2/loader/__init__.py b/nf2/loader/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/nf2/loader/analytical.py b/nf2/loader/analytical.py
new file mode 100644
index 0000000..c747604
--- /dev/null
+++ b/nf2/loader/analytical.py
@@ -0,0 +1,151 @@
+import os
+
+import numpy as np
+import wandb
+from matplotlib import pyplot as plt
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, RandomSampler
+
+from nf2.data.analytical_field import get_analytic_b_field
+from nf2.data.dataset import BatchesDataset, RandomCoordinateDataset, CubeDataset
+from nf2.data.loader import prep_b_data
+
+
+class AnalyticDataModule(LightningDataModule):
+
+    def __init__(self, case, height, spatial_norm, b_norm, work_directory, batch_size={"boundary": 1e4, "random": 2e4},
+                 iterations=1e5, num_workers=None, boundary={"type": "full"}, **kwargs):
+        super().__init__()
+
+        # data parameters
+        self.spatial_norm = spatial_norm
+        self.height = height
+        self.b_norm = b_norm
+        self.meta_data = None
+
+        # train parameters
+        self.iterations = int(iterations)
+        self.num_workers = num_workers if num_workers is not None else os.cpu_count()
+
+        os.makedirs(work_directory, exist_ok=True)
+
+        tau_surfaces = boundary['tau_surfaces'] if boundary['type'] == 'tau' else None
+        if case == 1:
+            b_cube = get_analytic_b_field(n=1, m=1, l=0.3, psi=np.pi / 4, tau_surfaces=tau_surfaces)
+        elif case == 2:
+            b_cube = get_analytic_b_field(n=1, m=1, l=0.3, psi=np.pi * 0.15, tau_surfaces=tau_surfaces,
+                                          resolution=[80, 80, 72])
+        else:
+            raise Exception(f'Invalid CASE {case}. Available cases are: [1, 2]')
+
+        for i in range(b_cube.shape[2]):
+            fig, axs = plt.subplots(1, 3, figsize=(6, 2))
+            min_max = np.abs(b_cube[..., i, 0]).max()
+            axs[0].imshow(b_cube[..., i, 0].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
+            axs[1].imshow(b_cube[..., i, 1].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
+            axs[2].imshow(b_cube[..., i, 2].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
+            wandb.log({"Overview": fig})
+            plt.close('all')
+
+        if boundary['type'] == "full":
+            coord_cube = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1)
+            values = [b_cube[:, :, :1].reshape((-1, 3)), b_cube[:, :, -1:].reshape((-1, 3)),
+                      b_cube[:, :1, :].reshape((-1, 3)), b_cube[:, -1:, :].reshape((-1, 3)),
+                      b_cube[:1, :, :].reshape((-1, 3)), b_cube[-1:, :, :].reshape((-1, 3)), ]
+            coords = [coord_cube[:, :, :1].reshape((-1, 3)), coord_cube[:, :, -1:].reshape((-1, 3)),
+                      coord_cube[:, :1, :].reshape((-1, 3)), coord_cube[:, -1:, :].reshape((-1, 3)),
+                      coord_cube[:1, :, :].reshape((-1, 3)), coord_cube[-1:, :, :].reshape((-1, 3)), ]
+            #
+            coords = np.concatenate(coords).astype(np.float32)
+            values = np.concatenate(values).astype(np.float32)
+        elif boundary['type'] == "potential":
+            b_cube = b_cube[:, :, 0]
+            coords, values, err = prep_b_data(b_cube, np.zeros_like(b_cube), height,
+                                              potential_boundary=True,
+                                              potential_strides=boundary['strides'])
+        elif boundary['type'] == "open":
+            # load dataset
+            coords = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1).astype(np.float32)
+            height_slices = boundary['height_slices'] if 'height_slices' in boundary else 0
+            b_cube = b_cube[:, :, height_slices]
+            coords = coords[:, :, height_slices]
+            # flatten data
+            coords = coords.reshape((-1, 3)).astype(np.float32)
+            values = b_cube.reshape((-1, 3)).astype(np.float32)
+        elif boundary['type'] == "tau":
+
+            # load dataset
+            coords = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1).astype(np.float32)
+            ranges = np.zeros((*coords.shape[:3], 2), dtype=np.float32)
+
+            height_mapping = {"z": [h / 2 for h in tau_surfaces],
+                              "z_min": [0] * len(tau_surfaces),
+                              "z_max": [h for h in tau_surfaces]}
+            self.height_mapping = height_mapping
+            for i, (z, z_min, z_max) in enumerate(
+                    zip(height_mapping["z"], height_mapping["z_min"], height_mapping["z_max"])):
+                ranges[:, :, i, 0] = z_min
+                ranges[:, :, i, 1] = z_max
+                coords[:, :, i, 2] = z
+            #
+            if boundary['use_LOS']:
+                b_cube[:, :, 1:, 0] = None
+                b_cube[:, :, 1:, 1] = None
+            # flatten data
+            coords = coords.reshape((-1, 3)).astype(np.float32)
+            values = b_cube.reshape((-1, 3)).astype(np.float32)
+            ranges = ranges.reshape((-1, 2)).astype(np.float32)
+        else:
+            raise Exception(f'Invalid boundary condition: {boundary}. Available options: ["full", "potential", "open"]')
+
+        # normalize data
+        values = values / b_norm
+        # apply spatial normalization
+        coords = coords / spatial_norm
+
+        cube_shape = [*b_cube.shape[:2], height]
+        self.cube_shape = cube_shape
+
+        # prep dataset
+        # shuffle data
+        r = np.random.permutation(coords.shape[0])
+        coords = coords[r]
+        values = values[r]
+        # store data to disk
+        coords_npy_path = os.path.join(work_directory, 'coords.npy')
+        np.save(coords_npy_path, coords)
+        values_npy_path = os.path.join(work_directory, 'values.npy')
+        np.save(values_npy_path, values)
+        # create data loaders
+        batches_path = {'coords': coords_npy_path, 'values': values_npy_path}
+        if boundary['type'] == "tau":  # add ranges
+            ranges /= spatial_norm
+            ranges = ranges[r]
+            ranges_npy_path = os.path.join(work_directory, 'ranges.npy')
+            np.save(ranges_npy_path, ranges)
+            batches_path['height_ranges'] = ranges_npy_path
+
+        boundary_batch_size = int(batch_size['boundary']) if isinstance(batch_size, dict) else int(batch_size)
+        random_batch_size = int(batch_size['random']) if isinstance(batch_size, dict) else int(batch_size)
+
+        self.dataset = BatchesDataset(batches_path, boundary_batch_size)
+        self.random_dataset = RandomCoordinateDataset(cube_shape, spatial_norm, random_batch_size)
+        self.cube_dataset = CubeDataset(cube_shape, spatial_norm, batch_size=boundary_batch_size)
+        self.batches_path = batches_path
+
+    def clear(self):
+        [os.remove(f) for f in self.batches_path.values()]
+
+    def train_dataloader(self):
+        data_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                 sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
+        random_loader = DataLoader(self.random_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                   sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
+        return {'boundary': data_loader, 'random': random_loader}
+
+    def val_dataloader(self):
+        cube_loader = DataLoader(self.cube_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                 shuffle=False)
+        boundary_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                     shuffle=False)
+        return [cube_loader, boundary_loader]
diff --git a/nf2/loader/base.py b/nf2/loader/base.py
new file mode 100644
index 0000000..8839d1e
--- /dev/null
+++ b/nf2/loader/base.py
@@ -0,0 +1,332 @@
+import os
+import uuid
+
+import numpy as np
+import wandb
+from astropy.nddata import block_reduce
+from matplotlib import pyplot as plt
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader, RandomSampler
+
+from nf2.data.dataset import CubeDataset, RandomCoordinateDataset, BatchesDataset
+from nf2.data.loader import load_potential_field_data
+from nf2.loader.util import _plot_B, _plot_B_error, _plot_los_trv_azi
+
+
+class TensorsDataset(BatchesDataset):
+
+    def __init__(self, tensors, work_directory, filter_nans=True, shuffle=True, ds_name=None, **kwargs):
+        # filter nan entries
+        nan_mask = np.any([np.all(np.isnan(t), axis=tuple(range(1, t.ndim))) for t in tensors.values()], axis=0)
+        if nan_mask.sum() > 0 and filter_nans:
+            print(f'Filtering {nan_mask.sum()} nan entries')
+            tensors = {k: v[~nan_mask] for k, v in tensors.items()}
+
+        # shuffle data
+        if shuffle:
+            r = np.random.permutation(list(tensors.values())[0].shape[0])
+            tensors = {k: v[r] for k, v in tensors.items()}
+
+        ds_name = uuid.uuid4() if ds_name is None else ds_name
+        batches_paths = {}
+        for k, v in tensors.items():
+            coords_npy_path = os.path.join(work_directory, f'{ds_name}_{k}.npy')
+            np.save(coords_npy_path, v.astype(np.float32))
+            batches_paths[k] = coords_npy_path
+
+        super().__init__(batches_paths, **kwargs)
+
+
+class BaseDataModule(LightningDataModule):
+
+    def __init__(self, training_datasets, validation_datasets, module_config, num_workers=None, iterations=None,
+                 **kwargs):
+        super().__init__()
+        self.training_datasets = training_datasets
+        self.validation_datasets = validation_datasets
+        self.datasets = {**self.training_datasets, **self.validation_datasets}
+
+        self.config = module_config
+        self.validation_dataset_mapping = {i: name for i, name in enumerate(self.validation_datasets.keys())}
+        self.num_workers = num_workers if num_workers is not None else os.cpu_count()
+        self.iterations = iterations
+
+    def clear(self):
+        [ds.clear() for ds in self.datasets.values() if isinstance(ds, BatchesDataset)]
+
+    def train_dataloader(self):
+        datasets = self.training_datasets
+
+        # data loader with fixed number of iterations
+        if self.iterations is not None:
+            loaders = {}
+            for i, (name, dataset) in enumerate(datasets.items()):
+                sampler = RandomSampler(dataset, replacement=True, num_samples=int(self.iterations))
+                loaders[name] = DataLoader(dataset, batch_size=None, num_workers=self.num_workers,
+                                           pin_memory=True, sampler=sampler)
+            return loaders
+
+        # data loader with iterations based on the largest dataset
+        ref_idx = np.argmax([len(ds) for ds in datasets.values()])
+        ref_dataset_name, ref_dataset = list(datasets.items())[ref_idx]
+        loaders = {ref_dataset_name: DataLoader(ref_dataset, batch_size=None, num_workers=self.num_workers,
+                                                pin_memory=True, shuffle=True)}
+        for i, (name, dataset) in enumerate(datasets.items()):
+            if i == ref_idx:
+                continue  # reference dataset already added
+            sampler = RandomSampler(dataset, replacement=True, num_samples=len(ref_dataset))
+            loaders[name] = DataLoader(dataset, batch_size=None, num_workers=self.num_workers,
+                                       pin_memory=True, sampler=sampler)
+        return loaders
+
+    def val_dataloader(self):
+        datasets = self.validation_datasets
+        loaders = []
+        for dataset in datasets.values():
+            loader = DataLoader(dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                shuffle=False)
+            loaders.append(loader)
+        return loaders
+
+
+class MapDataset(TensorsDataset):
+
+    def __init__(self, b, b_err=None, coords=None,
+                 G_per_dB=2500, Mm_per_pixel=0.36, Mm_per_ds=.36 * 320,
+                 bin=1, height_mapping=None, plot=True, los_trv_azi=False,
+                 wcs=None, **kwargs):
+        self.ds_per_pixel = (Mm_per_pixel * bin) / Mm_per_ds
+
+        # binning
+        b = block_reduce(b, (bin, bin, 1), np.mean)
+        # normalize
+        if los_trv_azi:
+            b[..., :2] /= G_per_dB
+        else:
+            b /= G_per_dB
+
+        if coords is None:
+            coords = np.stack(np.mgrid[:b.shape[0], :b.shape[1], :1], -1).astype(np.float32) * self.ds_per_pixel
+            coords = coords[:, :, 0, :]
+        else:
+            coords = coords / Mm_per_ds
+
+        self.coord_range = np.array([[coords[..., 0].min(), coords[..., 0].max()],
+                                     [coords[..., 1].min(), coords[..., 1].max()]])
+
+        self.cube_shape = coords.shape[:-1]
+        self.los_trv_azi = los_trv_azi
+        self.height_mapping = height_mapping
+        self.wcs = wcs
+
+        tensors = {'b_true': b}
+
+        if height_mapping is not None:
+            z = height_mapping['z']
+            z_min = height_mapping['z_min'] if 'z_min' in height_mapping else 0
+            z_max = height_mapping['z_max'] if 'z_max' in height_mapping else 0
+
+            coords[..., 2] = z / Mm_per_ds
+            ranges = np.zeros((*self.cube_shape, 2), dtype=np.float32)
+            ranges[..., 0] = z_min / Mm_per_ds
+            ranges[..., 1] = z_max / Mm_per_ds
+            tensors['height_range'] = ranges
+
+        tensors['coords'] = coords
+
+        if b_err is not None:
+            b_err = block_reduce(b_err, (bin, bin, 1), np.mean)
+            if los_trv_azi:
+                b_err[..., :2] /= G_per_dB
+            else:
+                b_err /= G_per_dB
+            tensors['b_err'] = b_err
+
+            if plot and not los_trv_azi:
+                _plot_B_error(b * G_per_dB, b_err * G_per_dB, coords * Mm_per_ds)
+        else:
+            if plot and los_trv_azi:
+                b_plot = np.copy(b)
+                b_plot[..., :2] *= G_per_dB
+                _plot_los_trv_azi(b_plot, coords * Mm_per_ds)
+            if plot and not los_trv_azi:
+                _plot_B(b * G_per_dB, coords * Mm_per_ds)
+
+
+        tensors = {k: v.reshape((-1, *v.shape[2:])).astype(np.float32) for k, v in tensors.items()}
+
+        super().__init__(tensors, **kwargs)
+
+
+class SlicesDataModule(BaseDataModule):
+
+    def __init__(self, b_slices, work_directory, height=160, spatial_norm=160, b_norm=2500,
+                 batch_size={"boundary": 1e4, "random": 2e4},
+                 iterations=1e5, num_workers=None,
+                 error_slices=None, height_mapping={'z': [0]}, boundary={"type": "open"},
+                 validation_strides=1,
+                 meta_data=None, plot_overview=True, Mm_per_pixel=None, buffer=None,
+                 **kwargs):
+        super().__init__()
+
+        # data parameters
+        self.spatial_norm = spatial_norm
+        self.height = height
+        self.b_norm = b_norm
+        self.height_mapping = height_mapping
+        self.meta_data = meta_data
+        self.Mm_per_pixel = Mm_per_pixel
+
+        # train parameters
+        self.iterations = int(iterations)
+        self.num_workers = num_workers if num_workers is not None else os.cpu_count()
+
+        os.makedirs(work_directory, exist_ok=True)
+
+        self.b_slices = b_slices
+
+        if plot_overview:
+            for i in range(b_slices.shape[2]):
+                fig, axs = plt.subplots(1, 3, figsize=(12, 4))
+                axs[0].imshow(b_slices[..., i, 0].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+                axs[1].imshow(b_slices[..., i, 1].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+                axs[2].imshow(b_slices[..., i, 2].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+                wandb.log({"Overview": fig})
+                plt.close('all')
+
+        # load dataset
+        assert len(height_mapping['z']) == b_slices.shape[
+            2], 'Invalid height mapping configuration: z must have the same length as the number of slices'
+        coords = np.stack(np.mgrid[:b_slices.shape[0], :b_slices.shape[1], :b_slices.shape[2]], -1).astype(np.float32)
+        for i, h in enumerate(height_mapping['z']):
+            coords[:, :, i, 2] = h
+        ranges = np.zeros((*coords.shape[:-1], 2))
+        use_height_range = 'z_max' in height_mapping
+        if use_height_range:
+            z1 = height_mapping['z_max']
+            # set to lower boundary if not specified
+            z0 = height_mapping['z_min'] if 'z_min' in height_mapping else np.zeros_like(z1)
+            assert len(z0) == len(z1) == len(height_mapping['z']), \
+                'Invalid height mapping configuration: z_min, z_max and z must have the same length'
+            for i, (h_min, h_max) in enumerate(zip(z0, z1)):
+                ranges[:, :, i, 0] = h_min
+                ranges[:, :, i, 1] = h_max
+        # flatten data
+        coords = coords.reshape((-1, 3)).astype(np.float32)
+        values = b_slices.reshape((-1, 3)).astype(np.float32)
+        ranges = ranges.reshape((-1, 2)).astype(np.float32)
+        errors = error_slices.reshape((-1, 3)).astype(np.float32) if error_slices is not None else np.zeros_like(values)
+
+        # filter nan entries
+        nan_mask = np.all(np.isnan(values), -1)
+        if nan_mask.sum() > 0:
+            print(f'Filtering {nan_mask.sum()} nan entries')
+            coords = coords[~nan_mask]
+            values = values[~nan_mask]
+            ranges = ranges[~nan_mask]
+            errors = errors[~nan_mask]
+
+        if boundary['type'] == 'potential':
+            b_bottom = b_slices[:, :, 0]
+            b_bottom = np.nan_to_num(b_bottom, nan=0)  # replace nans of mosaic data
+            pf_coords, pf_errors, pf_values = load_potential_field_data(b_bottom, height, boundary['strides'],
+                                                                        progress=True)
+            #
+            pf_ranges = np.zeros((*pf_coords.shape[:-1], 2), dtype=np.float32)
+            pf_ranges[:, 0] = pf_coords[:, 2]
+            pf_ranges[:, 1] = pf_coords[:, 2]
+            # concatenate pf data points
+            coords = np.concatenate([pf_coords, coords])
+            values = np.concatenate([pf_values, values])
+            ranges = np.concatenate([pf_ranges, ranges])
+            errors = np.concatenate([pf_errors, errors])
+        elif boundary['type'] == 'potential_top':
+            b_bottom = b_slices[:, :, 0]
+            b_bottom = np.nan_to_num(b_bottom, nan=0)  # replace nans of mosaic data
+            pf_coords, pf_errors, pf_values = load_potential_field_data(b_bottom, height, boundary['strides'],
+                                                                        only_top=True, pf_error=0.1, progress=True)
+            # log upper boundary
+            fig, axs = plt.subplots(1, 3, figsize=(12, 4))
+            pf_b_map = pf_values.reshape(b_bottom.shape)
+            axs[0].imshow(pf_b_map[..., 0].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+            axs[1].imshow(pf_b_map[..., 1].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+            axs[2].imshow(pf_b_map[..., 2].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
+            wandb.log({"Overview": fig})
+            plt.close('all')
+            #
+            pf_ranges = np.zeros((*pf_coords.shape[:-1], 2), dtype=np.float32)
+            pf_ranges[:, 0] = height / 2
+            pf_ranges[:, 1] = height
+            # concatenate pf data points
+            coords = np.concatenate([pf_coords, coords])
+            values = np.concatenate([pf_values, values])
+            ranges = np.concatenate([pf_ranges, ranges])
+            errors = np.concatenate([pf_errors, errors])
+        elif boundary['type'] == 'open':
+            pass
+        else:
+            raise ValueError('Unknown boundary type')
+
+        # normalize data
+        values = values / b_norm
+        errors = errors / b_norm
+        # apply spatial normalization
+        coords = coords / spatial_norm
+        ranges = ranges / spatial_norm
+
+        cube_shape = [*b_slices.shape[:-2], height]
+        self.cube_shape = cube_shape
+
+        # prep dataset
+        # shuffle data
+        r = np.random.permutation(coords.shape[0])
+        coords = coords[r]
+        values = values[r]
+        ranges = ranges[r]
+        errors = errors[r]
+        # store data to disk
+        coords_npy_path = os.path.join(work_directory, 'coords.npy')
+        np.save(coords_npy_path, coords)
+        values_npy_path = os.path.join(work_directory, 'values.npy')
+        np.save(values_npy_path, values)
+        batches_path = {'coords': coords_npy_path,
+                        'values': values_npy_path, }
+
+        # add height ranges if provided
+        if use_height_range:
+            ranges_npy_path = os.path.join(work_directory, 'ranges.npy')
+            np.save(ranges_npy_path, ranges)
+            batches_path['height_ranges'] = ranges_npy_path
+
+        # add error ranges if provided
+        if error_slices is not None:
+            err_npy_path = os.path.join(work_directory, 'errors.npy')
+            np.save(err_npy_path, errors)
+            batches_path['errors'] = err_npy_path
+
+        boundary_batch_size = int(batch_size['boundary']) if isinstance(batch_size, dict) else int(batch_size)
+        random_batch_size = int(batch_size['random']) if isinstance(batch_size, dict) else int(batch_size)
+
+        # create data loaders
+        self.dataset = BatchesDataset(batches_path, boundary_batch_size)
+        self.random_dataset = RandomCoordinateDataset(cube_shape, spatial_norm, random_batch_size, buffer=buffer)
+        self.cube_dataset = CubeDataset(cube_shape, spatial_norm, batch_size=boundary_batch_size,
+                                        strides=validation_strides)
+        self.batches_path = batches_path
+
+    def clear(self):
+        [os.remove(f) for f in self.batches_path.values()]
+
+    def train_dataloader(self):
+        data_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                 sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
+        random_loader = DataLoader(self.random_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                   sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
+        return {'boundary': data_loader, 'random': random_loader}
+
+    def val_dataloader(self):
+        cube_loader = DataLoader(self.cube_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                 shuffle=False)
+        boundary_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
+                                     shuffle=False)
+        return [boundary_loader, cube_loader]
diff --git a/nf2/loader/fits.py b/nf2/loader/fits.py
new file mode 100644
index 0000000..650b00a
--- /dev/null
+++ b/nf2/loader/fits.py
@@ -0,0 +1,250 @@
+import glob
+import os
+from copy import copy
+
+import numpy as np
+from astropy import units as u
+from astropy.io import fits
+from astropy.nddata import block_reduce
+from sunpy.map import Map
+
+from nf2.data.dataset import RandomCoordinateDataset, CubeDataset, SlicesDataset
+from nf2.data.loader import load_potential_field_data
+from nf2.loader.base import TensorsDataset, BaseDataModule, MapDataset
+
+
+class FITSDataModule(BaseDataModule):
+
+    def __init__(self, slices, work_directory, boundary_config=None,
+                 random_config=None,
+                 Mm_per_ds=.36 * 320, G_per_dB=2500, max_height=100, validation_batch_size=2 ** 15, log_shape=False,
+                 **kwargs):
+        # wrap data if only one slice is provided
+        slices = slices if isinstance(slices, list) else [slices]
+        # boundary dataset
+        slice_datasets = []
+        for config in slices:
+            boundary_dataset = self.init_boundary_dataset(**config,
+                                                          Mm_per_ds=Mm_per_ds, G_per_dB=G_per_dB,
+                                                          work_directory=work_directory, )
+            slice_datasets.append(boundary_dataset)
+
+        bottom_boundary_dataset = slice_datasets[0]
+
+        # random sampling dataset
+        coord_range = bottom_boundary_dataset.coord_range
+        z_range = np.array([[0, max_height / Mm_per_ds]])
+        coord_range = np.concatenate([coord_range, z_range], axis=0)
+        random_config = random_config if random_config is not None else {}
+        random_dataset = RandomCoordinateDataset(coord_range, **random_config)
+
+        ds_per_pixel = bottom_boundary_dataset.ds_per_pixel
+
+        if log_shape:
+            print(f'EXTRAPOLATING CUBE:')
+            # pretty plot cube range
+            print(f'x: {coord_range[0, 0] * Mm_per_ds:.2f} - {coord_range[0, 1] * Mm_per_ds:.2f} Mm')
+            print(f'y: {coord_range[1, 0] * Mm_per_ds:.2f} - {coord_range[1, 1] * Mm_per_ds:.2f} Mm')
+            print(f'z: {coord_range[2, 0] * Mm_per_ds:.2f} - {coord_range[2, 1] * Mm_per_ds:.2f} Mm')
+            print('------------------')
+            print(f'x: {coord_range[0, 0] / ds_per_pixel :.2f} - {coord_range[0, 1] / ds_per_pixel:.2f} pixel')
+            print(f'y: {coord_range[1, 0] / ds_per_pixel:.2f} - {coord_range[1, 1] / ds_per_pixel:.2f} pixel')
+            print(f'z: {coord_range[2, 0] / ds_per_pixel:.2f} - {coord_range[2, 1] / ds_per_pixel:.2f} pixel')
+            print('------------------')
+            print(f'x: {coord_range[0, 0]:.2f} - {coord_range[0, 1]:.2f} ds')
+            print(f'y: {coord_range[1, 0]:.2f} - {coord_range[1, 1]:.2f} ds')
+            print(f'z: {coord_range[2, 0]:.2f} - {coord_range[2, 1]:.2f} ds')
+
+        training_datasets = {}
+        for i, dataset in enumerate(slice_datasets):
+            training_datasets[f'boundary_{i + 1:02d}'] = dataset
+        training_datasets['random'] = random_dataset
+
+        # top and side boundaries
+        boundary_config = boundary_config if boundary_config is not None else \
+            {'type': 'potential', 'strides': 4,
+             'file': slices[0]['fits_path']['Br'] if 'Br' in slices[0]['fits_path'] else slices[0]['fits_path'][
+                 'B_los'],
+             'Mm_per_pixel': ds_per_pixel * Mm_per_ds}
+        if boundary_config['type'] == 'potential':
+            bz = fits.getdata(boundary_config['file']).transpose()
+            bz = np.nan_to_num(bz, nan=0)  # replace nans with 0
+            potential_dataset = PotentialBoundaryDataset(bz=bz,
+                                                         height_pixel=max_height / boundary_config['Mm_per_pixel'],
+                                                         ds_per_pixel=ds_per_pixel, G_per_dB=G_per_dB,
+                                                         work_directory=work_directory,
+                                                         strides=boundary_config['strides'])
+            training_datasets['potential'] = potential_dataset
+
+        # validation datasets
+        cube_dataset = CubeDataset(coord_range, batch_size=validation_batch_size)
+        validation_slice_datasets = []
+        for slice_config in slices:
+            validation_boundary_dataset = self.init_boundary_dataset(**slice_config,
+                                                                     Mm_per_ds=Mm_per_ds, G_per_dB=G_per_dB,
+                                                                     shuffle=False, filter_nans=False,
+                                                                     work_directory=work_directory,
+                                                                     batch_size=validation_batch_size, plot=False, )
+            validation_slice_datasets.append(validation_boundary_dataset)
+        validation_slices_dataset = SlicesDataset(coord_range, ds_per_pixel, n_slices=10,
+                                                  batch_size=validation_batch_size)
+
+        validation_datasets = {'cube': cube_dataset, 'slices': validation_slices_dataset}
+        for i, dataset in enumerate(validation_slice_datasets):
+            validation_datasets[f'validation_boundary_{i + 1:02d}'] = dataset
+
+        config = {'type': 'cartesian',
+                  'Mm_per_ds': Mm_per_ds, 'G_per_dB': G_per_dB, 'max_height': max_height,
+                  'coord_range': [], 'ds_per_pixel': [], 'height_mapping': [], 'wcs': []}
+        for ds in slice_datasets:
+            config['coord_range'].append(ds.coord_range)
+            config['ds_per_pixel'].append(ds.ds_per_pixel)
+            config['height_mapping'].append(ds.height_mapping)
+            config['wcs'].append(ds.wcs)
+
+        super().__init__(training_datasets, validation_datasets, config, **kwargs)
+
+    def init_boundary_dataset(self, type='fits', **kwargs):
+        if type == 'fits':
+            return FITSDataset(**kwargs)
+        elif type == 'los_trv_azi':
+            return LosTrvAziFITSDataset(**kwargs)
+        elif type == 'sharp':
+            return SHARPDataset(**kwargs)
+        else:
+            raise ValueError(f'Unknown boundary type: {type}')
+
+
+class FITSDataset(MapDataset):
+
+    def __init__(self, fits_path, error_path=None, coords_path=None,
+                 bin=1, slice=None, **kwargs):
+
+        file_p = fits_path['Bp']
+        file_t = fits_path['Bt']
+        file_r = fits_path['Br']
+        p_map, t_map, r_map = Map(file_p), Map(file_t), Map(file_r)
+        p_map = process_map(p_map, slice, bin)
+        t_map = process_map(t_map, slice, bin)
+        r_map = process_map(r_map, slice, bin)
+
+        b = np.stack([p_map.data, -t_map.data, r_map.data]).transpose()
+        wcs = r_map.wcs
+
+        if error_path is not None:
+            if isinstance(error_path, str):
+                file_p_err = sorted(glob.glob(os.path.join(error_path, '*Bp_err.fits')))[0]  # x
+                file_t_err = sorted(glob.glob(os.path.join(error_path, '*Bt_err.fits')))[0]  # y
+                file_r_err = sorted(glob.glob(os.path.join(error_path, '*Br_err.fits')))[0]  # z
+            else:
+                file_p_err = error_path['Bp_err']
+                file_t_err = error_path['Bt_err']
+                file_r_err = error_path['Br_err']
+            p_error_map, t_error_map, r_error_map = Map(file_p_err), Map(file_t_err), Map(file_r_err)
+            p_error_map = process_map(p_error_map, slice, bin)
+            t_error_map = process_map(t_error_map, slice, bin)
+            r_error_map = process_map(r_error_map, slice, bin)
+            b_err = np.stack([p_error_map.data, t_error_map.data, r_error_map.data]).transpose()
+        else:
+            b_err = None
+
+        if coords_path is not None:
+            coords_x, coords_y, coords_z = Map(coords_path['x']), Map(coords_path['y']), Map(coords_path['z'])
+            coords_x = process_map(coords_x, slice, bin)
+            coords_y = process_map(coords_y, slice, bin)
+            coords_z = process_map(coords_z, slice, bin)
+            coords = np.stack([coords_x.data, coords_y.data, coords_z.data]).transpose()
+        else:
+            coords = None
+
+        super().__init__(b=b, b_err=b_err, coords=coords, wcs=wcs, **kwargs)
+
+
+class LosTrvAziFITSDataset(MapDataset):
+
+    def __init__(self, fits_path,
+                 bin=1, slice=None, load_map=True, **kwargs):
+        file_los = fits_path['B_los']
+        file_trv = fits_path['B_trv']
+        file_azi = fits_path['B_azi']
+
+        if load_map:
+            los_map, trv_map, azi_map = Map(file_los), Map(file_trv), Map(file_azi)
+            los_map = process_map(los_map, slice, bin)
+            trv_map = process_map(trv_map, slice, bin)
+            azi_map = process_map(azi_map, slice, bin)
+            los_data = los_map.data
+            trv_data = trv_map.data
+            azi_data = azi_map.data
+            wcs = los_map.wcs
+        else:
+            los_data = fits.getdata(file_los)
+            trv_data = fits.getdata(file_trv)
+            azi_data = fits.getdata(file_azi)
+            if slice:
+                los_data = los_data[slice[0]:slice[1], slice[2]:slice[3]]
+                trv_data = trv_data[slice[0]:slice[1], slice[2]:slice[3]]
+                azi_data = azi_data[slice[0]:slice[1], slice[2]:slice[3]]
+            if bin > 1:
+                los_data = block_reduce(los_data, (bin, bin), func=np.mean)
+                trv_data = block_reduce(trv_data, (bin, bin), func=np.mean)
+                azi_data = block_reduce(azi_data, (bin, bin), func=np.mean)
+            wcs = None
+
+        b = np.stack([los_data, trv_data, np.pi - azi_data]).transpose()
+
+        super().__init__(b=b, wcs=wcs, los_trv_azi=True, **kwargs)
+
+
+class FITSSeriesDataModule(FITSDataModule):
+
+    def __init__(self, fits_paths, error_paths=None, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        self.fits_paths = copy(fits_paths)
+        self.error_paths = copy(error_paths) if error_paths is not None else [None] * len(fits_paths)
+
+        self.current_id = os.path.basename(self.fits_paths[0]['Br']).split('.')[-3]
+
+        self.initialized = True  # only required for first iteration
+        super().__init__({'fits_path': self.fits_paths[0], 'error_path': self.error_paths[0]}, *args, **kwargs)
+
+    def train_dataloader(self):
+        if not self.initialized:
+            # update ID
+            self.current_id = os.path.basename(self.fits_paths[0]['Br']).split('.')[-3]
+            # re-initialize
+            super().__init__(slices={'fits_path': self.fits_paths[0], 'error_path': self.error_paths[0]},
+                             *self.args, **self.kwargs)
+        # continue with next file in list
+        del self.fits_paths[0]
+        del self.error_paths[0]
+        self.initialized = False
+        return super().train_dataloader()
+
+
+class SHARPDataset(FITSDataset):
+
+    def __init__(self, **kwargs):
+        super().__init__(Mm_per_pixel=.36, **kwargs)
+
+
+class PotentialBoundaryDataset(TensorsDataset):
+
+    def __init__(self, bz, height_pixel, ds_per_pixel, G_per_dB, strides=2, batch_size=2 ** 10, **kwargs):
+        coords, b_err, b = load_potential_field_data(bz, height_pixel, strides,
+                                                     progress=False)
+        coords = coords * ds_per_pixel
+        b_err = b_err / G_per_dB
+        b = b / G_per_dB
+
+        super().__init__({'b_true': b, 'b_err': b_err, 'coords': coords}, batch_size=batch_size, **kwargs)
+
+
+def process_map(map, slice, bin):
+    if slice:
+        map = map.submap(bottom_left=u.Quantity((slice[0], slice[2]), u.pixel),
+                         top_right=u.Quantity((slice[1], slice[3]), u.pixel))
+    if bin > 1:
+        map = map.superpixel(u.Quantity((bin, bin), u.pixel), func=np.mean)
+    return map
diff --git a/nf2/loader/general.py b/nf2/loader/general.py
new file mode 100644
index 0000000..83b1a18
--- /dev/null
+++ b/nf2/loader/general.py
@@ -0,0 +1,22 @@
+import numpy as np
+from astropy.nddata import block_reduce
+
+from nf2.loader.base import SlicesDataModule
+
+
+class NumpyDataModule(SlicesDataModule):
+
+    def __init__(self, data_path, slices=None, bin=1, use_bz=False, components=False, *args, **kwargs):
+        b_slices = np.load(data_path)
+        if slices:
+            b_slices = b_slices[:, :, slices]
+        if bin > 1:
+            b_slices = block_reduce(b_slices, (bin, bin, 1, 1), np.mean)
+        if use_bz:
+            b_slices[:, :, 1:, 0] = None
+            b_slices[:, :, 1:, 1] = None
+        if components:
+            for i, c in enumerate(components):
+                filter = [i for i in [0, 1, 2] if i not in c]
+                b_slices[:, :, i, filter] = None
+        super().__init__(b_slices, *args, **kwargs)
diff --git a/nf2/loader/muram.py b/nf2/loader/muram.py
new file mode 100644
index 0000000..37195e0
--- /dev/null
+++ b/nf2/loader/muram.py
@@ -0,0 +1,216 @@
+import os
+
+import numpy as np
+from astropy.nddata import block_reduce
+
+from nf2.data.dataset import RandomCoordinateDataset, CubeDataset, SlicesDataset
+from nf2.data.util import img_to_los_trv_azi
+from nf2.loader.base import MapDataset, BaseDataModule
+from nf2.loader.fits import PotentialBoundaryDataset
+
+from astropy import units as u
+
+class MURaMDataModule(BaseDataModule):
+
+    def __init__(self, slices, work_directory, boundary_config=None, random_config=None,
+                 Mm_per_ds=.36 * 320, G_per_dB=2500, max_height=100, validation_batch_size=2 ** 15, log_shape=False,
+                 **kwargs):
+        # boundary dataset
+        slice_datasets = []
+        for slice_config in slices:
+            muram_dataset = MURaMDataset(**slice_config, G_per_dB=G_per_dB, Mm_per_ds=Mm_per_ds, work_directory=work_directory)
+            slice_datasets.append(muram_dataset)
+
+        bottom_boundary_dataset = slice_datasets[0]
+
+        # random sampling dataset
+        coord_range = bottom_boundary_dataset.coord_range
+        z_range = np.array([[0, max_height / Mm_per_ds]])
+        coord_range = np.concatenate([coord_range, z_range], axis=0)
+        random_config = random_config if random_config is not None else {}
+        random_dataset = RandomCoordinateDataset(coord_range, **random_config)
+
+        ds_per_pixel = bottom_boundary_dataset.ds_per_pixel
+
+        if log_shape:
+            print(f'EXTRAPOLATING CUBE:')
+            # pretty plot cube range
+            print(f'x: {coord_range[0, 0] * Mm_per_ds:.2f} - {coord_range[0, 1] * Mm_per_ds:.2f} Mm')
+            print(f'y: {coord_range[1, 0] * Mm_per_ds:.2f} - {coord_range[1, 1] * Mm_per_ds:.2f} Mm')
+            print(f'z: {coord_range[2, 0] * Mm_per_ds:.2f} - {coord_range[2, 1] * Mm_per_ds:.2f} Mm')
+            print('------------------')
+            print(f'x: {coord_range[0, 0] / ds_per_pixel :.2f} - {coord_range[0, 1] / ds_per_pixel:.2f} pixel')
+            print(f'y: {coord_range[1, 0] / ds_per_pixel:.2f} - {coord_range[1, 1] / ds_per_pixel:.2f} pixel')
+            print(f'z: {coord_range[2, 0] / ds_per_pixel:.2f} - {coord_range[2, 1] / ds_per_pixel:.2f} pixel')
+            print('------------------')
+            print(f'x: {coord_range[0, 0]:.2f} - {coord_range[0, 1]:.2f} ds')
+            print(f'y: {coord_range[1, 0]:.2f} - {coord_range[1, 1]:.2f} ds')
+            print(f'z: {coord_range[2, 0]:.2f} - {coord_range[2, 1]:.2f} ds')
+
+        training_datasets = {}
+        for i, dataset in enumerate(slice_datasets):
+            training_datasets[f'boundary_{i + 1:02d}'] = dataset
+        training_datasets['random'] = random_dataset
+
+        # top and side boundaries
+        boundary_config = boundary_config if boundary_config is not None else {'type': 'potential', 'strides': 8}
+        if boundary_config['type'] == 'potential':
+            sl, Nvar, shape, time = read_muram_slice(slices[0]['data_path'])
+            bz = sl[5, :, :] * np.sqrt(4 * np.pi)
+            potential_dataset = PotentialBoundaryDataset(bz=bz, height_pixel=max_height / (ds_per_pixel * Mm_per_ds),
+                                                         ds_per_pixel=ds_per_pixel, G_per_dB=G_per_dB,
+                                                         work_directory=work_directory,
+                                                         strides=boundary_config['strides'])
+            training_datasets['potential'] = potential_dataset
+
+        # validation datasets
+        cube_dataset = CubeDataset(coord_range, batch_size=validation_batch_size)
+
+        validation_slice_datasets = []
+        for slice_config in slices:
+            slice_config['batch_size'] = validation_batch_size # override batch size
+            muram_dataset = MURaMDataset(**slice_config, G_per_dB=G_per_dB, Mm_per_ds=Mm_per_ds, work_directory=work_directory,
+                                         shuffle=False, filter_nans=False, plot=False)
+            validation_slice_datasets.append(muram_dataset)
+        validation_slices_dataset = SlicesDataset(coord_range, ds_per_pixel, n_slices=10,
+                                                  batch_size=validation_batch_size)
+
+        validation_datasets = {'cube': cube_dataset, 'slices': validation_slices_dataset}
+        for i, dataset in enumerate(validation_slice_datasets):
+            validation_datasets[f'validation_boundary_{i + 1:02d}'] = dataset
+
+        config = {'type': 'cartesian',
+                  'Mm_per_ds': Mm_per_ds, 'G_per_dB': G_per_dB, 'max_height': max_height,
+                  'coord_range': [], 'ds_per_pixel': [], 'height_mapping': []}
+        for ds in slice_datasets:
+            config['coord_range'].append(ds.coord_range)
+            config['ds_per_pixel'].append(ds.ds_per_pixel)
+            config['height_mapping'].append(ds.height_mapping)
+
+        super().__init__(training_datasets, validation_datasets, config, **kwargs)
+
+
+muram_variables = {'Bz': {'id': 'result_prim_5', 'unit': u.Gauss},
+             'By': {'id': 'result_prim_7', 'unit': u.Gauss},
+             'Bx': {'id': 'result_prim_6', 'unit': u.Gauss},
+             'vx': {'id': 'result_prim_1', 'unit': u.cm / u.s},
+             'vy': {'id': 'result_prim_2', 'unit': u.cm / u.s},
+             'vz': {'id': 'result_prim_3', 'unit': u.cm / u.s},
+             'tau': {'id': 'tau', 'unit': 1},
+                   }
+
+class MURaMDataset(MapDataset):
+
+    def __init__(self, data_path, *args, **kwargs):
+        sl, Nvar, shape, time = read_muram_slice(data_path)
+
+        bz = sl[5, :, :] * np.sqrt(4 * np.pi)
+        bx = sl[6, :, :] * np.sqrt(4 * np.pi)
+        by = sl[7, :, :] * np.sqrt(4 * np.pi)
+
+        b = np.stack([bx, by, bz], axis=-1)
+
+        super().__init__(b, Mm_per_pixel=0.192, *args, **kwargs)
+
+class MURaMLosTrvAziDataset(MapDataset):
+
+    def __init__(self, data_path, *args, **kwargs):
+        sl, Nvar, shape, time = read_muram_slice(data_path)
+
+        bz = sl[5, :, :] * np.sqrt(4 * np.pi)
+        bx = sl[6, :, :] * np.sqrt(4 * np.pi)
+        by = sl[7, :, :] * np.sqrt(4 * np.pi)
+
+        b = np.stack([bx, by, bz], axis=-1)
+        b = img_to_los_trv_azi(b, f=np)
+
+        super().__init__(b, Mm_per_pixel=0.192, los_trv_azi=True, *args, **kwargs)
+
+class MURaMSnapshot():
+
+    def __init__(self, source_path, iteration):
+        header_file = os.path.join(source_path, f'Header.{iteration:06d}')
+        assert os.path.exists(header_file), f'Header file {header_file} does not exist'
+
+        # load shape information
+        header = np.loadtxt(header_file, dtype=np.float32)
+        shape = header[0:3].astype(int)
+        ds = header[3:6] * u.cm / u.pix  # cm; pixel size
+        # range = (shape * u.pix) * ds  # cm; domain size
+        time = header[6] * u.s
+        # (5, 6, 7) --> (z, x, y)
+
+        files = {k: os.path.join(source_path, f"{v['id']}.{iteration:06d}") for k, v in muram_variables.items()}
+        # (x, y, z) --> (z, y, x) --> (y, z, x)
+        data = {k: np.memmap(f, mode='r', dtype=np.float32, shape=tuple(shape[::-1])).transpose(1, 0, 2)
+                for k, f in files.items()}
+
+        for k, v in data.items():
+            setattr(self, k, v)
+        self.time = time
+        self.shape = (shape[1], shape[2], shape[0])
+        self.ds = (ds[1], ds[2], ds[0])
+
+    @property
+    def B(self):
+        return np.stack([self.Bx, self.By, self.Bz], axis=-1) * np.sqrt(4 * np.pi)
+
+    def v(self):
+        return np.stack([self.vx, self.vy, self.vz], axis=-1)
+
+    def load_cube(self, resolution=0.192 * u.Mm / u.pix, height=100 * u.Mm):
+
+        b = self.B
+        tau = self.tau
+
+        # integer division
+        assert resolution % self.ds[0] == 0, f'resolution {resolution} must be a multiple of {self.ds[0]}'
+        assert resolution % self.ds[1] == 0, f'resolution {resolution} must be a multiple of {self.ds[1]}'
+        assert resolution % self.ds[2] == 0, f'resolution {resolution} must be a multiple of {self.ds[2]}'
+        x_binning = resolution // self.ds[0]
+        y_binning = resolution // self.ds[1]
+        z_binning = resolution // self.ds[2]
+
+        b = block_reduce(b, (x_binning, y_binning, z_binning, 1), np.mean)
+        tau = block_reduce(tau, (x_binning, y_binning, z_binning), np.mean)
+
+        pix_height = np.argmin(np.abs(tau - 1), axis=2) * u.pix
+        base_height_pix = pix_height.mean()
+
+        min_height = int(base_height_pix.to_value(u.pix))
+        max_height = int((base_height_pix + height / resolution).to_value(u.pix))
+        b = b[:, :, min_height:max_height]
+        tau = tau[:, :, min_height:max_height]
+
+        return {'B': b, 'tau': tau}
+
+    def load_base(self, resolution=0.192 * u.Mm / u.pix, height=100 * u.Mm, base_height=180):
+        b = self.B[:, :, base_height:]
+        tau = self.tau[:, :, base_height:]
+        v = np.stack([self.vx, self.vy, self.vz], axis=-1)
+
+        # integer division
+        assert resolution % self.ds[0] == 0, f'resolution {resolution} must be a multiple of {self.ds[0]}'
+        assert resolution % self.ds[1] == 0, f'resolution {resolution} must be a multiple of {self.ds[1]}'
+        assert resolution % self.ds[2] == 0, f'resolution {resolution} must be a multiple of {self.ds[2]}'
+        x_binning = resolution // self.ds[0]
+        y_binning = resolution // self.ds[1]
+        z_binning = resolution // self.ds[2]
+
+        b = block_reduce(b, (x_binning, y_binning, z_binning, 1), np.mean)
+        tau = block_reduce(tau, (x_binning, y_binning, z_binning), np.mean)
+
+
+        max_height = int((height / resolution).to_value(u.pix))
+        b = b[:, :, :max_height]
+        tau = tau[:, :, :max_height]
+
+        return {'B': b, 'tau': tau}
+
+def read_muram_slice(filepath):
+    data = np.fromfile(filepath, dtype=np.float32)
+    Nvar = data[0].astype(int)
+    shape = tuple(data[1:3].astype(int))
+    time = data[3]
+    slice = data[4:].reshape([Nvar, shape[1], shape[0]]).swapaxes(1, 2)
+    return slice, Nvar, shape, time
diff --git a/nf2/loader/spherical.py b/nf2/loader/spherical.py
new file mode 100644
index 0000000..d42d96e
--- /dev/null
+++ b/nf2/loader/spherical.py
@@ -0,0 +1,330 @@
+import os
+from collections import OrderedDict
+from copy import copy
+
+import numpy as np
+import pfsspy
+import wandb
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
+
+from nf2.data.dataset import RandomSphericalCoordinateDataset, SphereDataset, SphereSlicesDataset
+from nf2.data.util import spherical_to_cartesian, vector_spherical_to_cartesian, cartesian_to_spherical_matrix
+from nf2.loader.base import BaseDataModule, TensorsDataset
+
+
+class SphericalSliceDataset(TensorsDataset):
+
+    def __init__(self, b, coords, spherical_coords, G_per_dB, Mm_per_ds,
+                 b_err=None, transform=None,
+                 plot_overview=True, strides=1, **kwargs):
+        if plot_overview:
+            self._plot(b, coords, spherical_coords)
+        if strides > 1:
+            b = b[::strides, ::strides]
+            coords = coords[::strides, ::strides]
+            b_err = b_err[::strides, ::strides] if b_err is not None else None
+            transform = transform[::strides, ::strides] if transform is not None else None
+        self.cube_shape = b.shape[:-1]
+        # flatten data
+        b = np.concatenate([f.reshape((-1, 3)) for f in b]).astype(np.float32)
+        coords = np.concatenate([c.reshape((-1, 3)) for c in coords]).astype(np.float32)
+        if b_err is not None:
+            b_err = np.concatenate([e.reshape((-1, 3)) for e in b_err]).astype(np.float32)
+        if transform is not None:
+            transform = np.concatenate([t.reshape((-1, 3, 3)) for t in transform]).astype(np.float32)
+
+        # normalize data
+        b /= G_per_dB
+        b_err = b_err / G_per_dB if b_err is not None else None
+        coords = coords * (1 * u.solRad).to_value(u.Mm) / Mm_per_ds
+
+        tensors = {'coords': coords,
+                   'b_true': b, }
+        if transform is not None:
+            tensors['transform'] = transform
+        if b_err is not None:
+            tensors['b_err'] = b_err
+
+        super().__init__(tensors, **kwargs)
+
+    def _plot(self, b, coords, spherical_coords):
+        b_min_max = np.nanmax(np.abs(b))
+        fig, axs = plt.subplots(3, 1, figsize=(8, 8))
+        im = axs[0].imshow(b[..., 0].transpose(), vmin=-b_min_max, vmax=b_min_max, cmap='gray', origin='lower')
+        divider = make_axes_locatable(axs[0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax)
+        im = axs[1].imshow(b[..., 1].transpose(), vmin=-b_min_max, vmax=b_min_max, cmap='gray', origin='lower')
+        divider = make_axes_locatable(axs[1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax)
+        im = axs[2].imshow(b[..., 2].transpose(), vmin=-b_min_max, vmax=b_min_max, cmap='gray', origin='lower')
+        divider = make_axes_locatable(axs[2])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        fig.colorbar(im, cax=cax)
+        wandb.log({"Overview - B": fig})
+        plt.close('all')
+
+        fig, axs = plt.subplots(3, 1, figsize=(8, 8))
+        im = axs[0].imshow(coords[..., 0].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[0])
+        im = axs[1].imshow(coords[..., 1].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[1])
+        im = axs[2].imshow(coords[..., 2].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[2])
+        wandb.log({"Coordinates": fig})
+        plt.close('all')
+
+        fig, axs = plt.subplots(3, 1, figsize=(8, 8))
+        im = axs[0].imshow(spherical_coords[..., 0].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[0])
+        im = axs[1].imshow(spherical_coords[..., 1].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[1])
+        im = axs[2].imshow(spherical_coords[..., 2].transpose(), origin='lower')
+        fig.colorbar(im, ax=axs[2])
+        wandb.log({"Spherical Coordinates": fig})
+        plt.close('all')
+
+
+class SphericalMapDataset(SphericalSliceDataset):
+
+    def __init__(self, files, mask_configs=[], **kwargs):
+        # load maps
+        r_map = Map(files['Br'])
+        t_map = Map(files['Bt'])
+        p_map = Map(files['Bp'])
+
+        # load coordinates
+        spherical_coords = all_coordinates_from_map(r_map).transform_to(frames.HeliographicCarrington)
+        r = spherical_coords.radius
+        r = r * u.solRad if r.unit == u.dimensionless_unscaled else r
+        spherical_coords = np.stack([
+            r.to_value(u.solRad),
+            np.pi / 2 - spherical_coords.lat.to_value(u.rad),
+            spherical_coords.lon.to_value(u.rad),
+        ]).transpose()
+        cartesian_coords = spherical_to_cartesian(spherical_coords)
+
+        # load data and transform matrix
+        b_spherical = np.stack([r_map.data, t_map.data, p_map.data]).transpose()
+        b_cartesian = vector_spherical_to_cartesian(b_spherical, spherical_coords)
+        transform = cartesian_to_spherical_matrix(spherical_coords)
+
+        # load error maps
+        if 'Br_err' in files and 'Bt_err' in files and 'Bp_err' in files:
+            r_error_map = Map(files['Br_err'])
+            t_error_map = Map(files['Bt_err'])
+            p_error_map = Map(files['Bp_err'])
+            b_error_spherical = np.stack([r_error_map.data,
+                                          t_error_map.data,
+                                          p_error_map.data]).transpose()
+        else:
+            b_error_spherical = None
+
+        # apply masking
+        mask_configs = mask_configs if isinstance(mask_configs, list) else [mask_configs]
+        for mask_config in mask_configs:
+            self._mask(b_cartesian, b_error_spherical, b_spherical, cartesian_coords, mask_config, r_map,
+                       spherical_coords,
+                       transform)
+        if np.any(np.isnan(cartesian_coords)):
+            # crop nan values
+            nan_coords = np.argwhere(~np.isnan(cartesian_coords).any(-1))
+            min_x = nan_coords[..., 0].min()
+            max_x = nan_coords[..., 0].max()
+            min_y = nan_coords[..., 1].min()
+            max_y = nan_coords[..., 1].max()
+            # truncate data
+            cartesian_coords = cartesian_coords[min_x:max_x, min_y:max_y]
+            spherical_coords = spherical_coords[min_x:max_x, min_y:max_y]
+            b_spherical = b_spherical[min_x:max_x, min_y:max_y]
+            if b_error_spherical is not None:
+                b_error_spherical = b_error_spherical[min_x:max_x, min_y:max_y]
+            if transform is not None:
+                transform = transform[min_x:max_x, min_y:max_y]
+
+        super().__init__(b=b_spherical, coords=cartesian_coords, spherical_coords=spherical_coords,
+                         b_err=b_error_spherical, transform=transform,
+                         **kwargs)
+
+    def _mask(self, b_cartesian, b_error_spherical, b_spherical, cartesian_coords, mask_config, r_map,
+              spherical_coords, transform):
+        m_type = mask_config['type']
+        if m_type == 'reference':
+            ref_map = Map(mask_config['file'])
+            ref_map.meta[
+                'date-obs'] = r_map.date.to_datetime().isoformat()  # use reference time to avoid temporal shift
+            reprojected_map = ref_map.reproject_to(r_map.wcs)
+            mask = ~np.isnan(reprojected_map.data).T
+        elif m_type == 'helioprojective':
+            ref_map = Map(mask_config['file'])
+            bottom_left = SkyCoord(mask_config['Tx'][0] * u.arcsec, mask_config['Ty'][0] * u.arcsec,
+                                   frame=ref_map.coordinate_frame)
+            top_right = SkyCoord(mask_config['Tx'][1] * u.arcsec, mask_config['Ty'][1] * u.arcsec,
+                                 frame=ref_map.coordinate_frame)
+            bottom_left = bottom_left.transform_to(frames.HeliographicCarrington)
+            top_right = top_right.transform_to(frames.HeliographicCarrington)
+            slice_lon = np.array([bottom_left.lon.to(u.rad).value, top_right.lon.to(u.rad).value])
+            slice_lat = np.array([bottom_left.lat.to(u.rad).value, top_right.lat.to(u.rad).value]) + np.pi / 2
+
+            mask = (spherical_coords[..., 1] > slice_lat[0]) & \
+                   (spherical_coords[..., 1] < slice_lat[1]) & \
+                   (spherical_coords[..., 2] > slice_lon[0]) & \
+                   (spherical_coords[..., 2] < slice_lon[1])
+        elif m_type == 'heliographic_carrington':
+            unit = u.Quantity(mask_config['unit']) if 'unit' in mask_config else u.rad
+            slice_lon = (mask_config['longitude_range'] * unit).to_value(u.rad)
+            slice_lat = (mask_config['latitude_range'] * unit).to_value(u.rad)
+
+            lat_mask = (spherical_coords[..., 1] > slice_lat[0]) & \
+                       (spherical_coords[..., 1] < slice_lat[1])
+            lon_mask = (spherical_coords[..., 2] > slice_lon[0]) & \
+                       (spherical_coords[..., 2] < slice_lon[1])
+            if slice_lon[1] > 2 * np.pi:
+                lon_mask = lon_mask | \
+                           ((spherical_coords[..., 2] > 0) & (spherical_coords[..., 2] < slice_lon[1] - 2 * np.pi))
+
+            mask = lat_mask & lon_mask
+        else:
+            raise NotImplementedError(f"Unknown mask type '{type}'. "
+                                      f"Please choose from: ['reference', 'helioprojective', 'heliographic_carrington']")
+
+        mask = ~mask if 'invert' in mask_config and mask_config['invert'] else mask
+
+        b_cartesian[mask] = np.nan
+        b_spherical[mask] = np.nan
+        if b_error_spherical is not None:
+            b_error_spherical[mask] = np.nan
+        cartesian_coords[mask] = np.nan
+        spherical_coords[mask] = np.nan
+        transform[mask] = np.nan
+
+
+class PFSSBoundaryDataset(SphericalSliceDataset):
+
+    def __init__(self, Br, radius_range, source_surface_height=2.5, resample=[360, 180], sampling_points=100, mask=None,
+                 insert=None, **kwargs):
+        height = radius_range[1]
+        assert source_surface_height >= height, 'Source surface height must be greater than height (set source_surface_height to >height)'
+
+        # load synoptic map
+        potential_r_map = Map(Br)
+        potential_r_map = potential_r_map.resample(resample * u.pix)
+
+        # insert additional data (e.g., HMI full disk maps)
+        insert = insert if insert is not None else []
+        insert = insert if isinstance(insert, list) else [insert]
+        for file_in in insert:
+            map_in = Map(file_in)
+            map_in = map_in.reproject_to(potential_r_map.wcs)
+            condition = ~np.isnan(map_in.data)
+            potential_r_map.data[condition] = map_in.data[condition]
+
+        # PFSS extrapolation
+        pfss_in = pfsspy.Input(potential_r_map, sampling_points, source_surface_height)
+        pfss_out = pfsspy.pfss(pfss_in)
+
+        # load B field
+        ref_coords = all_coordinates_from_map(potential_r_map)
+        spherical_coords = SkyCoord(lon=ref_coords.lon, lat=ref_coords.lat, radius=height * u.solRad,
+                                    frame=ref_coords.frame)
+        potential_shape = spherical_coords.shape  # required workaround for pfsspy spherical reshape
+        spherical_b = pfss_out.get_bvec(spherical_coords.reshape((-1,)))
+        spherical_b = spherical_b.reshape((*potential_shape, 3)).value
+        spherical_b = np.stack([spherical_b[..., 0],
+                                spherical_b[..., 1],
+                                spherical_b[..., 2]]).T
+
+        # load coordinates
+        spherical_coords = np.stack([
+            spherical_coords.radius.value,
+            np.pi / 2 - spherical_coords.lat.to_value(u.rad),
+            spherical_coords.lon.to(u.rad).value]).T
+
+        if mask is not None:
+            condition = (spherical_coords[..., 1] < mask['latitude_range'][0]) | \
+                        (spherical_coords[..., 1] > mask['latitude_range'][1]) | \
+                        (spherical_coords[..., 2] < mask['longitude_range'][0]) | \
+                        (spherical_coords[..., 2] > mask['longitude_range'][1])
+            spherical_b[condition] = np.nan
+            spherical_coords[condition] = np.nan
+
+        # convert to spherical coordinates
+        coords = spherical_to_cartesian(spherical_coords)
+        transform = cartesian_to_spherical_matrix(spherical_coords)
+
+        super().__init__(b=spherical_b, coords=coords, spherical_coords=spherical_coords, transform=transform, **kwargs)
+
+
+class SphericalDataModule(BaseDataModule):
+
+    def __init__(self, train_configs, validation_configs,
+                 max_radius=None, Mm_per_ds=.36 * 320, G_per_dB=None, work_directory=None,
+                 batch_size=4096, **kwargs):
+
+        self.ds_mapping = {'map': SphericalMapDataset,
+                           'pfss_boundary': PFSSBoundaryDataset,
+                           'random_spherical': RandomSphericalCoordinateDataset,
+                           'sphere': SphereDataset,
+                           'spherical_slices': SphereSlicesDataset}
+
+        # data parameters
+        self.max_radius = max_radius
+        self.G_per_dB = G_per_dB
+        self.Mm_per_ds = Mm_per_ds
+        self.cube_shape = [1, max_radius]
+        self.spatial_norm = 1 * u.solRad
+
+        # init boundary datasets
+        general_config = {'work_directory': work_directory, 'batch_size': batch_size, 'G_per_dB': G_per_dB,
+                          'radius_range': [1, max_radius], 'Mm_per_ds': Mm_per_ds}
+
+        config = {'type': 'spherical',
+                  'radius_range': [1, max_radius],
+                  'G_per_dB': G_per_dB,
+                  'Mm_per_ds': Mm_per_ds}
+
+        training_datasets = self.load_config(train_configs, general_config, prefix='train')
+        validation_datasets = self.load_config(validation_configs, general_config, prefix='validation')
+
+        super().__init__(training_datasets, validation_datasets, config, **kwargs)
+
+    def load_config(self, configs, general_config, prefix='train'):
+        datasets = OrderedDict()
+        for i, config in enumerate(configs):
+            c_type = config.pop('type')
+            c_name = config.pop('ds_id') if 'ds_id' in config else f'{prefix}_{c_type}_{i}'
+            config['ds_name'] = c_name
+            # update config with general config
+            for k, v in general_config.items():
+                if k not in config:
+                    config[k] = v
+            os.makedirs(config['work_directory'], exist_ok=True)
+            dataset = self.ds_mapping[c_type](**config)
+            datasets[c_name] = dataset
+        return datasets
+
+
+class SphericalSeriesDataModule(SphericalDataModule):
+    def __init__(self, full_disk_files, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+        self.full_disk_files = copy(full_disk_files)
+        self.current_id = os.path.basename(self.full_disk_files[0]['Br']).split('.')[-3]
+
+        super().__init__(full_disk_files=self.full_disk_files[0], *self.args, **self.kwargs)
+
+    def train_dataloader(self):
+        if len(self.full_disk_files) == 0:
+            return None
+        # re-initialize
+        print(f"Load next file: {os.path.basename(self.full_disk_files[0]['Br'])}")
+        self.current_files = self.full_disk_files[0]
+        super().__init__(full_disk_files=self.full_disk_files[0], *self.args, **self.kwargs)
+        del self.full_disk_files[0]
+        return super().train_dataloader()
diff --git a/nf2/loader/util.py b/nf2/loader/util.py
new file mode 100644
index 0000000..cf48d19
--- /dev/null
+++ b/nf2/loader/util.py
@@ -0,0 +1,133 @@
+import numpy as np
+import wandb
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def _plot_B(b, coords):
+    fig, axs = plt.subplots(3, 2, figsize=(10, 15))
+    im = axs[0, 0].imshow(b[..., 0].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[0, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 0].set_title('Bp')
+    im = axs[0, 1].imshow(coords[..., 0].T, origin='lower')
+    divider = make_axes_locatable(axs[0, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 1].set_title('Coordinate x [Mm]')
+    #
+    im = axs[1, 0].imshow(b[..., 1].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[1, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 0].set_title('Bt')
+    im = axs[1, 1].imshow(coords[..., 1].T, origin='lower')
+    divider = make_axes_locatable(axs[1, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 1].set_title('Coordinate y [Mm]')
+    #
+    im = axs[2, 0].imshow(b[..., 2].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[2, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 0].set_title('Br')
+    im = axs[2, 1].imshow(coords[..., 2].T, origin='lower')
+    divider = make_axes_locatable(axs[2, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 1].set_title('Coordinate z [Mm]')
+    wandb.log({'FITS data': wandb.Image(plt)})
+    plt.close(fig)
+
+def _plot_los_trv_azi(b, coords):
+    fig, axs = plt.subplots(3, 2, figsize=(10, 15))
+    im = axs[0, 0].imshow(b[..., 0].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[0, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 0].set_title('B_los')
+    im = axs[0, 1].imshow(coords[..., 0].T, origin='lower')
+    divider = make_axes_locatable(axs[0, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 1].set_title('Coordinate x [Mm]')
+    #
+    im = axs[1, 0].imshow(b[..., 1].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[1, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 0].set_title('B_trv')
+    im = axs[1, 1].imshow(coords[..., 1].T, origin='lower')
+    divider = make_axes_locatable(axs[1, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 1].set_title('Coordinate y [Mm]')
+    #
+    im = axs[2, 0].imshow(b[..., 2].T % (2 * np.pi), origin='lower', cmap='twilight', vmin=0, vmax=2 * np.pi)
+    divider = make_axes_locatable(axs[2, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 0].set_title('B_azi')
+    im = axs[2, 1].imshow(coords[..., 2].T, origin='lower')
+    divider = make_axes_locatable(axs[2, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 1].set_title('Coordinate z [Mm]')
+    wandb.log({'FITS data': wandb.Image(plt)})
+    plt.close(fig)
+
+
+def _plot_B_error(b, b_err, coords):
+    fig, axs = plt.subplots(3, 3, figsize=(15, 15))
+    im = axs[0, 0].imshow(b[..., 0].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[0, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 0].set_title('Bp')
+    im = axs[0, 1].imshow(b_err[..., 0].T, origin='lower', cmap='Reds', norm=LogNorm(vmin=1, vmax=1000))
+    divider = make_axes_locatable(axs[0, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 1].set_title('Bp error')
+    im = axs[0, 2].imshow(coords[..., 0].T, origin='lower')
+    divider = make_axes_locatable(axs[0, 2])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[0, 2].set_title('Coordinate x [Mm]')
+    #
+    im = axs[1, 0].imshow(b[..., 1].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[1, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 0].set_title('Bt')
+    im = axs[1, 1].imshow(b_err[..., 1].T, origin='lower', cmap='Reds', norm=LogNorm(vmin=1, vmax=1000))
+    divider = make_axes_locatable(axs[1, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 1].set_title('Bt error')
+    im = axs[1, 2].imshow(coords[..., 1].T, origin='lower')
+    divider = make_axes_locatable(axs[1, 2])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[1, 2].set_title('Coordinate y [Mm]')
+    #
+    im = axs[2, 0].imshow(b[..., 2].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    divider = make_axes_locatable(axs[2, 0])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 0].set_title('Br')
+    im = axs[2, 1].imshow(b_err[..., 2].T, origin='lower', cmap='Reds', norm=LogNorm(vmin=1, vmax=1000))
+    divider = make_axes_locatable(axs[2, 1])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 1].set_title('Br error')
+    im = axs[2, 2].imshow(coords[..., 2].T, origin='lower')
+    divider = make_axes_locatable(axs[2, 2])
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    plt.colorbar(im, cax=cax, orientation='vertical')
+    axs[2, 2].set_title('Coordinate z [Mm]')
+    wandb.log({'FITS data': wandb.Image(plt)})
+    plt.close(fig)
diff --git a/nf2/loader/vsm.py b/nf2/loader/vsm.py
new file mode 100644
index 0000000..aee8f07
--- /dev/null
+++ b/nf2/loader/vsm.py
@@ -0,0 +1,17 @@
+import numpy as np
+
+from nf2.loader.base import SlicesDataModule
+
+
+class VSMDataModule(SlicesDataModule):
+
+    def __init__(self, data_path, slices=None, *args, **kwargs):
+        dict_data = np.load(data_path, allow_pickle=True)
+        sharp_cube = dict_data.item().get('sharp')
+        vsm_cube = dict_data.item().get('vsm')
+        vsm_cube = np.stack([np.ones_like(vsm_cube) * np.nan, np.ones_like(vsm_cube) * np.nan, vsm_cube])
+        b_slices = np.stack([sharp_cube, vsm_cube], 1).T
+        if slices is not None:
+            b_slices = b_slices[:, :, slices]
+
+        super().__init__(b_slices, *args, **kwargs)
diff --git a/nf2/potential/potential_field.py b/nf2/potential/potential_field.py
index 3766ee3..4a2020b 100644
--- a/nf2/potential/potential_field.py
+++ b/nf2/potential/potential_field.py
@@ -35,6 +35,7 @@ def get_potential(b_n, height, batch_size=2048, strides=(1, 1, 1), progress=True
     # r = (x * y, 3); coords = (x*y*z, 3), c = (1, 3)
     # --> (x * y, x * y * z, 3) --> (x * y, x * y * z) --> (x * y * z)
     device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    batch_size = batch_size * torch.cuda.device_count() if torch.cuda.is_available() else batch_size
     with torch.no_grad():
         b_n = torch.tensor(b_n, dtype=torch.float32, )
         r_p = torch.tensor(r_p, dtype=torch.float32, )
@@ -47,9 +48,9 @@ def get_potential(b_n, height, batch_size=2048, strides=(1, 1, 1), progress=True
         for coord, in it:
             coord = coord.to(device)
             p_batch = model(coord)
-            potential += [p_batch]
+            potential += [p_batch.detach().cpu()]
 
-    potential = torch.cat(potential).view(cube_shape).cpu().numpy()
+    potential = torch.cat(potential).view(cube_shape).numpy()
     if strides != (1, 1, 1):
         potential = block_replicate(potential, strides, conserve_sum=False)
     return potential
diff --git a/nf2/train/callback.py b/nf2/train/callback.py
new file mode 100644
index 0000000..28518e1
--- /dev/null
+++ b/nf2/train/callback.py
@@ -0,0 +1,579 @@
+import numpy as np
+import wandb
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from pytorch_lightning import Callback
+import torch
+
+from nf2.data.util import cartesian_to_spherical, vector_cartesian_to_spherical, img_to_los_trv_azi, los_trv_azi_to_img
+from astropy import units as u
+
+class SphericalSlicesCallback(Callback):
+
+    def __init__(self, name, cube_shape, gauss_per_dB, Mm_per_ds):
+        self.name = name
+        self.cube_shape = cube_shape
+        self.gauss_per_dB = gauss_per_dB
+        self.Mm_per_ds = Mm_per_ds
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if self.name not in pl_module.validation_outputs:
+            return
+        outputs = pl_module.validation_outputs[self.name]
+        b = outputs['b'] * self.gauss_per_dB
+        j = outputs['j'] * self.gauss_per_dB / self.Mm_per_ds
+        coords = outputs['coords']
+
+        b_cube = b.reshape([*self.cube_shape, 3]).cpu().numpy()
+        j_cube = j.reshape([*self.cube_shape, 3]).cpu().numpy()
+        c_cube = coords.reshape([*self.cube_shape, 3]).cpu().numpy()
+
+        # transform to spherical coordinates
+        c_cube = cartesian_to_spherical(c_cube)
+        b_cube = vector_cartesian_to_spherical(b_cube, c_cube)
+
+        c_cube[..., 0] *= self.Mm_per_ds / (1 * u.solRad).to_value(u.Mm)
+
+        self.plot_b(b_cube, c_cube)
+        self.plot_current(j_cube, c_cube)
+
+    def plot_b(self, b, coords):
+        n_samples = b.shape[0]
+        fig, axs = plt.subplots(3, n_samples, figsize=(n_samples * 4, 12))
+        for i in range(3):
+            for j in range(n_samples):
+                v_min_max = np.max(np.abs(b[j, :, :]))
+                # extent = [coords[j, 0, 0, 2], coords[j, -1, -1, 2],
+                #           coords[j, 0, 0, 1], coords[j, -1, -1, 1]]
+                # extent = np.rad2deg(extent)
+                extent = None
+                height = coords[j, :, :, 0].mean()
+                im = axs[i, j].imshow(b[j, :, :, i], cmap='gray', vmin=-v_min_max, vmax=v_min_max,
+                                      origin='upper', extent=extent)
+                axs[i, j].set_xlabel('Longitude [deg]')
+                axs[i, j].set_ylabel('Latitude [deg]')
+                # add locatable colorbar
+                divider = make_axes_locatable(axs[i, j])
+                cax = divider.append_axes("right", size="5%", pad=0.05)
+                plt.colorbar(im, cax=cax, label='B [G]')
+                axs[i, j].set_title(f'{height:.02f} - $B_{["r", "t", "p"][i]}$')
+        fig.tight_layout()
+        wandb.log({f"{self.name} - B": fig})
+        plt.close('all')
+
+    def plot_current(self, j, coords):
+        j = (j ** 2).sum(-1) ** 0.5
+        n_samples = j.shape[0]
+        fig, axs = plt.subplots(1, n_samples, figsize=(n_samples * 4, 4))
+        for i in range(n_samples):
+            # extent = [coords[i, 0, 0, 2], coords[i, -1, -1, 2],
+            #           coords[i, 0, 0, 1], coords[i, -1, -1, 1]]
+            # extent = np.rad2deg(extent)
+            extent = None
+            height = coords[i, :, :, 0].mean()
+            im = axs[i].imshow(j[i, :, :], cmap='viridis', origin='upper', norm=LogNorm(), extent=extent)
+            axs[i].set_xlabel('Longitude [deg]')
+            axs[i].set_ylabel('Latitude [deg]')
+            # add locatable colorbar
+            divider = make_axes_locatable(axs[i])
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+            plt.colorbar(im, cax=cax, label='J [G/ds]')
+            axs[i].set_title(f'{height:.02f} - $|J|$')
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Current density": fig})
+        plt.close('all')
+        # plot integrated current density
+        j = np.sum(j, axis=0)
+        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+        im = ax.imshow(j, cmap='viridis', origin='upper', norm=LogNorm(), extent=extent)
+        # add locatable colorbar
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax)
+        #
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Integrated Current density": fig})
+        plt.close('all')
+
+
+class SlicesCallback(Callback):
+
+    def __init__(self, name, cube_shape, gauss_per_dB, Mm_per_ds):
+        self.name = name
+        self.cube_shape = cube_shape
+        self.gauss_per_dB = gauss_per_dB
+        self.Mm_per_ds = Mm_per_ds
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if self.name not in pl_module.validation_outputs:
+            return
+        outputs = pl_module.validation_outputs[self.name]
+        b = outputs['b'] * self.gauss_per_dB
+        j = outputs['j'] * self.gauss_per_dB / self.Mm_per_ds
+        coords = outputs['coords']
+
+        b_cube = b.reshape([*self.cube_shape, 3]).cpu().numpy()
+        j_cube = j.reshape([*self.cube_shape, 3]).cpu().numpy()
+        c_cube = coords.reshape([*self.cube_shape, 3]).cpu().numpy()
+
+        self.plot_b(b_cube, c_cube)
+        self.plot_current(j_cube, c_cube)
+        if 'p' in outputs:
+            p = outputs['p']
+            p = p.reshape([*self.cube_shape]).cpu().numpy()
+            self.plot_pressure(p, c_cube)
+
+    def plot_b(self, b, coords):
+        n_samples = b.shape[2]
+        fig, axs = plt.subplots(3, n_samples, figsize=(n_samples * 4, 12))
+        for i in range(3):
+            for j in range(n_samples):
+                v_min_max = np.max(np.abs(b[:, :, j]))
+                extent = [coords[0, 0, j, 0], coords[-1, -1, j, 0],
+                          coords[0, 0, j, 1], coords[-1, -1, j, 1]]
+                extent = np.array(extent) * self.Mm_per_ds
+                height = coords[:, :, j, 2].mean() * self.Mm_per_ds
+                im = axs[i, j].imshow(b[:, :, j, i].T, cmap='gray', vmin=-v_min_max, vmax=v_min_max,
+                                      origin='lower', extent=extent)
+                axs[i, j].set_xlabel('X [Mm]')
+                axs[i, j].set_ylabel('Y [Mm]')
+                # add locatable colorbar
+                divider = make_axes_locatable(axs[i, j])
+                cax = divider.append_axes("right", size="5%", pad=0.05)
+                plt.colorbar(im, cax=cax, label='B [G]')
+                axs[i, j].set_title(f'{height:.02f} - $B_{["x", "y", "z"][i]}$')
+        fig.tight_layout()
+        wandb.log({f"{self.name} - B": fig})
+        plt.close('all')
+
+    def plot_current(self, j, coords):
+        j = (j ** 2).sum(-1) ** 0.5
+        n_samples = j.shape[2]
+        fig, axs = plt.subplots(1, n_samples, figsize=(n_samples * 4, 4))
+        for i in range(n_samples):
+            extent = [coords[0, 0, i, 0], coords[-1, -1, i, 0],
+                      coords[0, 0, i, 1], coords[-1, -1, i, 1]]
+            extent = np.array(extent) * self.Mm_per_ds
+            height = coords[:, :, i, 2].mean() * self.Mm_per_ds
+            im = axs[i].imshow(j[:, :, i].T, cmap='plasma', origin='lower', norm=LogNorm(), extent=extent)
+            axs[i].set_xlabel('X [Mm]')
+            axs[i].set_ylabel('Y [Mm]')
+            # add locatable colorbar
+            divider = make_axes_locatable(axs[i])
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+            plt.colorbar(im, cax=cax, label='J [G/ds]')
+            axs[i].set_title(f'{height:.02f} - $|J|$')
+            axs[i].set_xlabel('X [Mm]')
+            axs[i].set_ylabel('Y [Mm]')
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Current density": fig})
+        plt.close('all')
+        # plot integrated current density
+        j = np.sum(j, axis=2)
+        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+        im = ax.imshow(j.T, cmap='plasma', origin='lower', norm=LogNorm(), extent=extent)
+        # add locatable colorbar
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax)
+        ax.set_xlabel('X [Mm]')
+        ax.set_ylabel('Y [Mm]')
+        #
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Integrated Current density": fig})
+        plt.close('all')
+
+    def plot_pressure(self, p, coords):
+        n_samples = p.shape[2]
+        fig, axs = plt.subplots(1, n_samples, figsize=(n_samples * 4, 4))
+        for i in range(n_samples):
+            extent = [coords[0, 0, i, 0], coords[-1, -1, i, 0],
+                      coords[0, 0, i, 1], coords[-1, -1, i, 1]]
+            extent = np.array(extent) * self.Mm_per_ds
+            height = coords[:, :, i, 2].mean() * self.Mm_per_ds
+            im = axs[i].imshow(p[:, :, i].T, cmap='plasma', origin='lower', norm=LogNorm(), extent=extent)
+            axs[i].set_xlabel('X [Mm]')
+            axs[i].set_ylabel('Y [Mm]')
+            # add locatable colorbar
+            divider = make_axes_locatable(axs[i])
+            cax = divider.append_axes("right", size="5%", pad=0.05)
+            plt.colorbar(im, cax=cax, label='P')
+            axs[i].set_title(f'{height:.02f}')
+            axs[i].set_xlabel('X [Mm]')
+            axs[i].set_ylabel('Y [Mm]')
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Plasma Pressure": fig})
+        plt.close('all')
+        # plot integrated current density
+        p = np.sum(p, axis=2)
+        fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+        im = ax.imshow(p.T, cmap='plasma', origin='lower', norm=LogNorm(), extent=extent)
+        # add locatable colorbar
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='P')
+        ax.set_xlabel('X [Mm]')
+        ax.set_ylabel('Y [Mm]')
+        #
+        fig.tight_layout()
+        wandb.log({f"{self.name} - Integrated Plasma Pressure": fig})
+        plt.close('all')
+
+
+
+
+class BoundaryCallback(Callback):
+
+    def __init__(self, validation_dataset_key, cube_shape, gauss_per_dB, Mm_per_ds):
+        self.validation_dataset_key = validation_dataset_key
+        self.cube_shape = cube_shape
+        self.gauss_per_dB = gauss_per_dB
+        self.Mm_per_ds = Mm_per_ds
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if self.validation_dataset_key not in pl_module.validation_outputs:
+            return
+
+        outputs = pl_module.validation_outputs[self.validation_dataset_key]
+        b = outputs['b']
+        b_true = outputs['b_true']
+
+        # apply transforms
+        if 'transform' in outputs:
+            transform = outputs['transform']
+            b = torch.einsum('ijk,ik->ij', transform, b)
+
+        b = b * self.gauss_per_dB
+        b_true = b_true * self.gauss_per_dB
+
+        # compute diff
+        b_diff = torch.abs(b - b_true)
+        b_diff = torch.nanmean(b_diff.pow(2).sum(-1).pow(0.5))
+        evaluation = {'b_diff': b_diff.detach()}
+
+        # compute diff error
+        if 'b_err' in outputs:
+            b_err = outputs['b_err'] * self.gauss_per_dB
+            b_diff_err = torch.clip(torch.abs(b - b_true) - b_err, 0)
+            b_diff_err = torch.nanmean(b_diff_err.pow(2).sum(-1).pow(0.5))
+            evaluation['b_diff_err'] = b_diff_err.detach()
+
+        wandb.log({"valid": {self.validation_dataset_key: evaluation}})
+
+        b = b.cpu().numpy().reshape([*self.cube_shape, 3])
+        b_true = b_true.cpu().numpy().reshape([*self.cube_shape, 3])
+
+        if 'original_coords' in outputs:
+            original_coords = outputs['original_coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            transformed_coords = outputs['coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            self.plot_b_coords(b, b_true, original_coords, transformed_coords)
+        else:
+            original_coords = outputs['coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            self.plot_b(b, b_true, original_coords)
+
+    def plot_b(self, b, b_true, original_coords):
+        extent = None
+
+        fig, axs = plt.subplots(3, 2, figsize=(8, 8))
+
+        b_norm = np.nanmax(np.abs(b_true))
+        b_norm = min(500, b_norm)
+
+        im = axs[0, 0].imshow(b[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[0, 1].imshow(b_true[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[1, 0].imshow(b[..., 1].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[1, 1].imshow(b_true[..., 1].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[2, 0].imshow(b[..., 2].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[2, 1].imshow(b_true[..., 2].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        fig.tight_layout()
+        wandb.log({f"{self.validation_dataset_key} - B": fig})
+        plt.close('all')
+
+    def plot_b_coords(self, b, b_true, original_coords, transformed_coords):
+        extent = [original_coords[..., 0].min(), original_coords[..., 0].max(),
+                  original_coords[..., 1].min(), original_coords[..., 1].max()]
+
+        fig, axs = plt.subplots(4, 2, figsize=(8, 8))
+
+        b_norm = np.nanmax(np.abs(b_true))
+        b_norm = min(500, b_norm)
+
+        im = axs[0, 0].imshow(b[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[0, 1].imshow(b_true[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[1, 0].imshow(b[..., 1].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[1, 1].imshow(b_true[..., 1].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[2, 0].imshow(b[..., 2].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[2, 1].imshow(b_true[..., 2].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+
+        im = axs[3, 0].imshow(transformed_coords[..., 2].T, cmap='inferno', origin='lower', vmin=0, extent=extent)
+        divider = make_axes_locatable(axs[3, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Z [Mm]')
+        axs[3, 0].set_title('Transformed z')
+        axs[3, 0].set_xlabel('X [Mm]')
+        axs[3, 0].set_ylabel('Y [Mm]')
+
+        im = axs[3, 1].imshow(original_coords[..., 2].T, cmap='inferno', origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[3, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Z [Mm]')
+        axs[3, 1].set_title('Original z')
+        axs[3, 1].set_xlabel('X [Mm]')
+        axs[3, 1].set_ylabel('Y [Mm]')
+
+        fig.tight_layout()
+        wandb.log({f"{self.validation_dataset_key} - B": fig})
+        plt.close('all')
+
+class LosTrvAziBoundaryCallback(Callback):
+
+    def __init__(self, validation_dataset_key, cube_shape, gauss_per_dB, Mm_per_ds):
+        self.validation_dataset_key = validation_dataset_key
+        self.cube_shape = cube_shape
+        self.gauss_per_dB = gauss_per_dB
+        self.Mm_per_ds = Mm_per_ds
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if self.validation_dataset_key not in pl_module.validation_outputs:
+            return
+
+        outputs = pl_module.validation_outputs[self.validation_dataset_key]
+        b = outputs['b']
+        b_true = outputs['b_true']
+
+        # apply transforms
+        if 'transform' in outputs:
+            transform = outputs['transform']
+            b = torch.einsum('ijk,ik->ij', transform, b)
+
+        b = img_to_los_trv_azi(b, f=torch)
+        b[..., 0:2] = b[..., 0:2] * self.gauss_per_dB
+        b_true[..., 0:2] = b_true[..., 0:2] * self.gauss_per_dB
+
+        b_xyz = los_trv_azi_to_img(b, f=torch)
+        b_true_xyz = los_trv_azi_to_img(b_true, f=torch)
+        # compute diff
+        b_diff = torch.abs(b_xyz - b_true_xyz)
+        b_diff = torch.nanmean(b_diff.pow(2).sum(-1).pow(0.5))
+        evaluation = {'b_diff': b_diff.detach()}
+        #
+        b_xyz = los_trv_azi_to_img(b, ambiguous=True, f=torch)
+        b_true_xyz = los_trv_azi_to_img(b_true, ambiguous=True, f=torch)
+        # compute diff
+        b_diff = torch.abs(b_xyz - b_true_xyz)
+        b_diff = torch.nanmean(b_diff.pow(2).sum(-1).pow(0.5))
+        evaluation['b_amb_diff'] = b_diff.detach()
+
+        wandb.log({"valid": {self.validation_dataset_key: evaluation}})
+
+        b = b.cpu().numpy().reshape([*self.cube_shape, 3])
+        b_true = b_true.cpu().numpy().reshape([*self.cube_shape, 3])
+
+        if 'original_coords' in outputs:
+            original_coords = outputs['original_coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            transformed_coords = outputs['coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            self.plot_b_coords(b, b_true, original_coords, transformed_coords)
+        else:
+            original_coords = outputs['coords'].cpu().numpy().reshape([*self.cube_shape, 3]) * self.Mm_per_ds
+            self.plot_b(b, b_true, original_coords)
+
+    def plot_b(self, b, b_true, original_coords):
+        extent = [original_coords[..., 0].min(), original_coords[..., 0].max(),
+                  original_coords[..., 1].min(), original_coords[..., 1].max()]
+
+        fig, axs = plt.subplots(3, 2, figsize=(8, 8))
+
+        b_norm = np.nanmax(np.abs(b_true[..., 0]))
+        b_norm = min(500, b_norm)
+
+        im = axs[0, 0].imshow(b[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[0, 0].set_title('B_los')
+
+        im = axs[0, 1].imshow(b_true[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[0, 1].set_title('B_los_true')
+
+        im = axs[1, 0].imshow(b[..., 1].T, cmap='gray', vmin=0, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[1, 0].set_title('B_trv')
+
+        im = axs[1, 1].imshow(b_true[..., 1].T, cmap='gray', vmin=0, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[1, 1].set_title('B_trv_true')
+
+        im = axs[2, 0].imshow(b[..., 2].T % (2 * np.pi), cmap='twilight', vmin=0, vmax=2 * np.pi, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[deg]')
+        axs[2, 0].set_title('B_azi')
+
+        im = axs[2, 1].imshow(b_true[..., 2].T % (2 * np.pi), cmap='twilight', vmin=0, vmax=2 * np.pi, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[deg]')
+        axs[2, 1].set_title('B_azi_true')
+
+        fig.tight_layout()
+        wandb.log({f"{self.validation_dataset_key} - B": fig})
+        plt.close('all')
+
+    def plot_b_coords(self, b, b_true, original_coords, transformed_coords):
+        extent = [original_coords[..., 0].min(), original_coords[..., 0].max(),
+                  original_coords[..., 1].min(), original_coords[..., 1].max()]
+
+        fig, axs = plt.subplots(4, 2, figsize=(8, 8))
+
+        b_norm = np.nanmax(np.abs(b_true[..., 0]))
+        b_norm = min(500, b_norm)
+
+        im = axs[0, 0].imshow(b[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[0, 0].set_title('B_los')
+
+        im = axs[0, 1].imshow(b_true[..., 0].T, cmap='gray', vmin=-b_norm, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[0, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[0, 1].set_title('B_los_true')
+
+        im = axs[1, 0].imshow(b[..., 1].T, cmap='gray', vmin=0, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[1, 0].set_title('B_trv')
+
+        im = axs[1, 1].imshow(b_true[..., 1].T, cmap='gray', vmin=0, vmax=b_norm, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[1, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[G]')
+        axs[1, 1].set_title('B_trv_true')
+
+        im = axs[2, 0].imshow(b[..., 2].T % (2 * np.pi), cmap='twilight', vmin=0, vmax=2 * np.pi, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[deg]')
+        axs[2, 0].set_title('B_azi')
+
+        im = axs[2, 1].imshow(b_true[..., 2].T % (2 * np.pi), cmap='twilight', vmin=0, vmax=2 * np.pi, origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[2, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='[deg]')
+        axs[2, 1].set_title('B_azi_true')
+
+        im = axs[3, 0].imshow(transformed_coords[..., 2].T, cmap='inferno', origin='lower', vmin=0, extent=extent)
+        divider = make_axes_locatable(axs[3, 0])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Z [Mm]')
+        axs[3, 0].set_title('Transformed z')
+        axs[3, 0].set_xlabel('X [Mm]')
+        axs[3, 0].set_ylabel('Y [Mm]')
+
+        im = axs[3, 1].imshow(original_coords[..., 2].T, cmap='inferno', origin='lower', extent=extent)
+        divider = make_axes_locatable(axs[3, 1])
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Z [Mm]')
+        axs[3, 1].set_title('Original z')
+        axs[3, 1].set_xlabel('X [Mm]')
+        axs[3, 1].set_ylabel('Y [Mm]')
+
+        fig.tight_layout()
+        wandb.log({f"{self.validation_dataset_key} - B": fig})
+        plt.close('all')
+
+
+class MetricsCallback(Callback):
+
+    def __init__(self, validation_dataset_key, gauss_per_dB, Mm_per_ds):
+        self.validation_dataset_key = validation_dataset_key
+        self.gauss_per_dB = gauss_per_dB
+        self.Mm_per_ds = Mm_per_ds
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        if self.validation_dataset_key not in pl_module.validation_outputs:
+            return
+        outputs = pl_module.validation_outputs[self.validation_dataset_key]
+
+        b = outputs['b'] * self.gauss_per_dB
+        j = outputs['j'] * self.gauss_per_dB / self.Mm_per_ds
+
+        div = outputs['div'] * self.gauss_per_dB / self.Mm_per_ds
+
+        norm = torch.norm(b, dim=-1) * torch.norm(j, dim=-1)
+        sigma = torch.norm(torch.cross(j, b, dim=-1), dim=-1) / norm
+        j_weight = torch.norm(j, dim=-1)
+        angle = (sigma * j_weight).sum() / j_weight.sum()
+        angle = torch.clip(angle, -1. + 1e-7, 1. - 1e-7)
+        theta_J = torch.arcsin(angle)
+        theta_J = torch.rad2deg(theta_J)
+
+        sigma_J = (torch.norm(torch.cross(j, b, dim=-1), dim=-1) / (torch.norm(b, dim=-1) + 1e-7)).sum() / (torch.norm(j, dim=-1).sum() + 1e-7)
+
+        b_norm = b.pow(2).sum(-1).pow(0.5) + 1e-7
+        div_loss = (div / b_norm).mean()
+
+        ff_loss = torch.cross(j, b, dim=-1).pow(2).sum(-1).pow(0.5) / b_norm
+        ff_loss = ff_loss.mean()
+
+        wandb.log({"valid": {"divergence": div_loss.cpu().numpy(),
+                             "force-free": ff_loss.cpu().numpy(),
+                             "sigma_J": sigma_J.cpu().numpy(),
+                             "theta_J": theta_J.cpu().numpy()}})
diff --git a/nf2/train/config_from_template.py b/nf2/train/config_from_template.py
new file mode 100644
index 0000000..43f8786
--- /dev/null
+++ b/nf2/train/config_from_template.py
@@ -0,0 +1,28 @@
+# create a new file based on a template file
+import argparse
+import os
+from string import Template
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Create a new config file for synoptic extrapolations based on a template file')
+    parser.add_argument('--template_file', help='template file')
+    parser.add_argument('--carrington_rotation', nargs='+', help='carrington rotations of the synoptic maps')
+    parser.add_argument('--output_dir', help='output directory', required=False, default='.')
+    parser.add_argument('--overwrite', help='overwrite existing files', action='store_true')
+    args = parser.parse_args()
+
+    with open(args.template_file, 'r') as template_file:
+        src = Template(template_file.read())
+        for id in args.carrington_rotation:
+            output_file = os.path.join(args.output_dir, f'{id}.json')
+            if os.path.exists(output_file) and not args.overwrite:
+                print(f'File {output_file} already exists. Skipping.')
+                continue
+            with open(output_file, 'w') as target_file:
+                target_file.write(src.substitute(carrington_rotation=id))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/train/data_loader.py b/nf2/train/data_loader.py
deleted file mode 100644
index 90b5cef..0000000
--- a/nf2/train/data_loader.py
+++ /dev/null
@@ -1,463 +0,0 @@
-import glob
-import os
-from copy import copy
-
-import numpy as np
-import wandb
-from astropy.io import fits
-from astropy.nddata import block_reduce
-from matplotlib import pyplot as plt
-from pytorch_lightning import LightningDataModule
-from sunpy.map import Map
-from torch.utils.data import DataLoader, RandomSampler
-
-from nf2.data.analytical_field import get_analytic_b_field
-from nf2.data.dataset import CubeDataset, RandomCoordinateDataset, BatchesDataset
-from nf2.data.loader import prep_b_data, load_potential_field_data
-from astropy import units as u
-
-class SlicesDataModule(LightningDataModule):
-
-    def __init__(self, b_slices, height, spatial_norm, b_norm, work_directory,
-                 batch_size={"boundary": 1e4, "random": 2e4},
-                 iterations=1e5, num_workers=None,
-                 error_slices=None, height_mapping={'z': [0]}, boundary={"type": "open"},
-                 validation_strides = 1,
-                 meta_data=None, plot_overview=True, Mm_per_pixel=None, buffer=None,
-                 **kwargs):
-        super().__init__()
-
-        # data parameters
-        self.spatial_norm = spatial_norm
-        self.height = height
-        self.b_norm = b_norm
-        self.height_mapping = height_mapping
-        self.meta_data = meta_data
-        self.Mm_per_pixel = Mm_per_pixel
-
-        # train parameters
-        self.iterations = int(iterations)
-        self.num_workers = num_workers if num_workers is not None else os.cpu_count()
-
-        os.makedirs(work_directory, exist_ok=True)
-
-        self.b_slices = b_slices
-
-        if plot_overview:
-            for i in range(b_slices.shape[2]):
-                fig, axs = plt.subplots(1, 3, figsize=(12, 4))
-                axs[0].imshow(b_slices[..., i, 0].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-                axs[1].imshow(b_slices[..., i, 1].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-                axs[2].imshow(b_slices[..., i, 2].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-                wandb.log({"Overview": fig})
-                plt.close('all')
-
-        # load dataset
-        assert len(height_mapping['z']) == b_slices.shape[2], 'Invalid height mapping configuration: z must have the same length as the number of slices'
-        coords = np.stack(np.mgrid[:b_slices.shape[0], :b_slices.shape[1], :b_slices.shape[2]], -1).astype(np.float32)
-        for i, h in enumerate(height_mapping['z']):
-            coords[:, :, i, 2] = h
-        ranges = np.zeros((*coords.shape[:-1], 2))
-        use_height_range = 'z_max' in height_mapping
-        if use_height_range:
-            z1 = height_mapping['z_max']
-            # set to lower boundary if not specified
-            z0 = height_mapping['z_min'] if 'z_min' in height_mapping else np.zeros_like(z1)
-            assert len(z0) == len(z1) == len(height_mapping['z']), \
-                'Invalid height mapping configuration: z_min, z_max and z must have the same length'
-            for i, (h_min, h_max) in enumerate(zip(z0, z1)):
-                ranges[:, :, i, 0] = h_min
-                ranges[:, :, i, 1] = h_max
-        # flatten data
-        coords = coords.reshape((-1, 3)).astype(np.float32)
-        values = b_slices.reshape((-1, 3)).astype(np.float32)
-        ranges = ranges.reshape((-1, 2)).astype(np.float32)
-        errors = error_slices.reshape((-1, 3)).astype(np.float32) if error_slices is not None else np.zeros_like(values)
-
-        # filter nan entries
-        nan_mask = np.all(np.isnan(values), -1)
-        if nan_mask.sum() > 0:
-            print(f'Filtering {nan_mask.sum()} nan entries')
-            coords = coords[~nan_mask]
-            values = values[~nan_mask]
-            ranges = ranges[~nan_mask]
-            errors = errors[~nan_mask]
-
-        if boundary['type'] == 'potential':
-            b_bottom = b_slices[:, :, 0]
-            b_bottom = np.nan_to_num(b_bottom, nan=0) # replace nans of mosaic data
-            pf_coords, pf_errors, pf_values = load_potential_field_data(b_bottom, height, boundary['strides'], progress=True)
-            #
-            pf_ranges = np.zeros((*pf_coords.shape[:-1], 2), dtype=np.float32)
-            pf_ranges[:, 0] = pf_coords[:, 2]
-            pf_ranges[:, 1] = pf_coords[:, 2]
-            # concatenate pf data points
-            coords = np.concatenate([pf_coords, coords])
-            values = np.concatenate([pf_values, values])
-            ranges = np.concatenate([pf_ranges, ranges])
-            errors = np.concatenate([pf_errors, errors])
-        elif boundary['type'] == 'potential_top':
-            b_bottom = b_slices[:, :, 0]
-            b_bottom = np.nan_to_num(b_bottom, nan=0) # replace nans of mosaic data
-            pf_coords, pf_errors, pf_values = load_potential_field_data(b_bottom, height, boundary['strides'], only_top=True, pf_error=0.1, progress=True)
-            # log upper boundary
-            fig, axs = plt.subplots(1, 3, figsize=(12, 4))
-            pf_b_map = pf_values.reshape(b_bottom.shape)
-            axs[0].imshow(pf_b_map[..., 0].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-            axs[1].imshow(pf_b_map[..., 1].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-            axs[2].imshow(pf_b_map[..., 2].transpose(), vmin=-b_norm, vmax=b_norm, cmap='gray', origin='lower')
-            wandb.log({"Overview": fig})
-            plt.close('all')
-            #
-            pf_ranges = np.zeros((*pf_coords.shape[:-1], 2), dtype=np.float32)
-            pf_ranges[:, 0] = height / 2
-            pf_ranges[:, 1] = height
-            # concatenate pf data points
-            coords = np.concatenate([pf_coords, coords])
-            values = np.concatenate([pf_values, values])
-            ranges = np.concatenate([pf_ranges, ranges])
-            errors = np.concatenate([pf_errors, errors])
-        elif boundary['type'] == 'open':
-            pass
-        else:
-            raise ValueError('Unknown boundary type')
-
-        # normalize data
-        values = values / b_norm
-        errors = errors / b_norm
-        # apply spatial normalization
-        coords = coords / spatial_norm
-        ranges = ranges / spatial_norm
-
-        cube_shape = [*b_slices.shape[:-2], height]
-        self.cube_shape = cube_shape
-
-        # prep dataset
-        # shuffle data
-        r = np.random.permutation(coords.shape[0])
-        coords = coords[r]
-        values = values[r]
-        ranges = ranges[r]
-        errors = errors[r]
-        # store data to disk
-        coords_npy_path = os.path.join(work_directory, 'coords.npy')
-        np.save(coords_npy_path, coords)
-        values_npy_path = os.path.join(work_directory, 'values.npy')
-        np.save(values_npy_path, values)
-        batches_path = {'coords': coords_npy_path,
-                        'values': values_npy_path, }
-
-        # add height ranges if provided
-        if use_height_range:
-            ranges_npy_path = os.path.join(work_directory, 'ranges.npy')
-            np.save(ranges_npy_path, ranges)
-            batches_path['height_ranges'] = ranges_npy_path
-
-        # add error ranges if provided
-        if error_slices is not None:
-            err_npy_path = os.path.join(work_directory, 'errors.npy')
-            np.save(err_npy_path, errors)
-            batches_path['errors'] = err_npy_path
-
-        boundary_batch_size = int(batch_size['boundary']) if isinstance(batch_size, dict) else int(batch_size)
-        random_batch_size = int(batch_size['random']) if isinstance(batch_size, dict) else int(batch_size)
-
-        # create data loaders
-        self.dataset = BatchesDataset(batches_path, boundary_batch_size)
-        self.random_dataset = RandomCoordinateDataset(cube_shape, spatial_norm, random_batch_size, buffer=buffer)
-        self.cube_dataset = CubeDataset(cube_shape, spatial_norm, batch_size=boundary_batch_size, strides=validation_strides)
-        self.batches_path = batches_path
-
-    def clear(self):
-        [os.remove(f) for f in self.batches_path.values()]
-
-    def train_dataloader(self):
-        data_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                 sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
-        random_loader = DataLoader(self.random_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                   sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
-        return {'boundary': data_loader, 'random': random_loader}
-
-    def val_dataloader(self):
-        cube_loader = DataLoader(self.cube_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                 shuffle=False)
-        boundary_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                     shuffle=False)
-        return [cube_loader, boundary_loader]
-
-
-class SHARPDataModule(SlicesDataModule):
-
-    def __init__(self, data_path, bin=1, slice=None, *args, **kwargs):
-        if isinstance(data_path, str):
-            hmi_p = sorted(glob.glob(os.path.join(data_path, '*Bp.fits')))[0]  # x
-            hmi_t = sorted(glob.glob(os.path.join(data_path, '*Bt.fits')))[0]  # y
-            hmi_r = sorted(glob.glob(os.path.join(data_path, '*Br.fits')))[0]  # z
-            err_p = sorted(glob.glob(os.path.join(data_path, '*Bp_err.fits')))[0]  # x
-            err_t = sorted(glob.glob(os.path.join(data_path, '*Bt_err.fits')))[0]  # y
-            err_r = sorted(glob.glob(os.path.join(data_path, '*Br_err.fits')))[0]  # z
-        else:
-            hmi_p, err_p, hmi_t, err_t, hmi_r, err_r = data_path
-        # laod maps
-        p_map, t_map, r_map = Map(hmi_p),Map(hmi_t),Map(hmi_r)
-        p_error_map, t_error_map, r_error_map = Map(err_p), Map(err_t), Map(err_r)
-
-        maps = [p_map, t_map, r_map, p_error_map, t_error_map, r_error_map]
-        if slice:
-            maps = [m.submap(bottom_left=u.Quantity((slice[0], slice[2]), u.pixel),
-                             top_right=u.Quantity((slice[1], slice[3]), u.pixel)) for m in maps]
-        if bin > 1:
-            maps = [m.superpixel(u.Quantity((bin, bin), u.pixel), func=np.mean) for m in maps]
-
-        hmi_data = np.stack([maps[0].data, -maps[1].data, maps[2].data]).transpose()
-        error_data = np.stack([maps[3].data, maps[4].data, maps[5].data]).transpose()
-
-        b_slices = hmi_data[:, :, None]
-        error_slices = error_data[:, :, None]
-        meta_data = maps[0].meta
-
-        super().__init__(b_slices, *args, error_slices=error_slices, meta_data=meta_data, **kwargs)
-
-class SOLISDataModule(SlicesDataModule):
-
-    def __init__(self, data_path, slices=None, *args, **kwargs):
-        dict_data = np.load(data_path, allow_pickle=True)
-        sharp_cube = dict_data.item().get('sharp')
-        vsm_cube = dict_data.item().get('vsm')
-        vsm_cube = np.stack([np.ones_like(vsm_cube) * np.nan, np.ones_like(vsm_cube) * np.nan, vsm_cube])
-        b_slices = np.stack([sharp_cube, vsm_cube], 1).T
-        if slices is not None:
-            b_slices = b_slices[:, :, slices]
-
-        super().__init__(b_slices, *args, **kwargs)
-
-
-class FITSDataModule(SlicesDataModule):
-
-    def __init__(self, data_paths, mask_path=None, bin=1, flip_sign=None, *args, **kwargs):
-        # check if single height layer
-        if isinstance(data_paths, dict):
-            data_paths = [data_paths]
-        # load all heights
-        b_slices = []
-        meta_data = None
-        for data_path in data_paths:
-            if isinstance(data_path, dict):
-                data_dict = data_path
-                b = np.stack([fits.getdata(data_dict['x']).T,
-                              fits.getdata(data_dict['y']).T,
-                              fits.getdata(data_dict['z']).T], -1)
-                meta_data = fits.getheader(data_dict['x'])
-            else:
-                b = fits.getdata(data_path).T
-                meta_data = fits.getheader(data_path)
-            b_slices += [b]
-        b_slices = np.stack(b_slices, 2)
-        if flip_sign is not None:
-            b_slices[..., flip_sign] *= -1  # -t component
-
-        if mask_path is not None:
-            mask = fits.getdata(mask_path).T
-            b_slices[mask == 0, :, :] = np.nan
-        if bin > 1:
-            b_slices = block_reduce(b_slices, (bin, bin, 1, 1), np.mean)
-        super().__init__(b_slices, meta_data=meta_data, *args, **kwargs)
-
-
-class SHARPSeriesDataModule(SHARPDataModule):
-
-    def __init__(self, file_paths, *args, **kwargs):
-        self.args = args
-        self.kwargs = kwargs
-        self.file_paths = copy(file_paths)
-
-        super().__init__(file_paths[0], *self.args, **self.kwargs)
-
-    def train_dataloader(self):
-        # re-initialize
-        super().__init__(self.file_paths[0], *self.args, **self.kwargs)
-        del self.file_paths[0] # continue with next file in list
-        return super().train_dataloader()
-
-
-class NumpyDataModule(SlicesDataModule):
-
-    def __init__(self, data_path, slices=None, bin=1, use_bz=False, components=False, *args, **kwargs):
-        b_slices = np.load(data_path)
-        if slices:
-            b_slices = b_slices[:, :, slices]
-        if bin > 1:
-            b_slices = block_reduce(b_slices, (bin, bin, 1, 1), np.mean)
-        if use_bz:
-            b_slices[:, :, 1:, 0] = None
-            b_slices[:, :, 1:, 1] = None
-        if components:
-            for i, c in enumerate(components):
-                filter = [i for i in [0, 1, 2] if i not in c]
-                b_slices[:, :, i, filter] = None
-        super().__init__(b_slices, *args, **kwargs)
-
-
-class AnalyticDataModule(LightningDataModule):
-
-    def __init__(self, case, height, spatial_norm, b_norm, work_directory, batch_size={"boundary": 1e4, "random": 2e4},
-                 iterations=1e5, num_workers=None, boundary={"type": "full"}, **kwargs):
-        super().__init__()
-
-        # data parameters
-        self.spatial_norm = spatial_norm
-        self.height = height
-        self.b_norm = b_norm
-        self.meta_data = None
-
-        # train parameters
-        self.iterations = int(iterations)
-        self.num_workers = num_workers if num_workers is not None else os.cpu_count()
-
-        os.makedirs(work_directory, exist_ok=True)
-
-        tau_surfaces = boundary['tau_surfaces'] if boundary['type'] == 'tau' else None
-        if case == 1:
-            b_cube = get_analytic_b_field(n=1, m=1, l=0.3, psi=np.pi / 4, tau_surfaces=tau_surfaces)
-        elif case == 2:
-            b_cube = get_analytic_b_field(n=1, m=1, l=0.3, psi=np.pi * 0.15, tau_surfaces=tau_surfaces,
-                                          resolution=[80, 80, 72])
-        else:
-            raise Exception(f'Invalid CASE {case}. Available cases are: [1, 2]')
-
-        for i in range(b_cube.shape[2]):
-            fig, axs = plt.subplots(1, 3, figsize=(6, 2))
-            min_max = np.abs(b_cube[..., i, 0]).max()
-            axs[0].imshow(b_cube[..., i, 0].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
-            axs[1].imshow(b_cube[..., i, 1].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
-            axs[2].imshow(b_cube[..., i, 2].transpose(), vmin=-min_max, vmax=min_max, cmap='gray', origin='lower')
-            wandb.log({"Overview": fig})
-            plt.close('all')
-
-        if boundary['type'] == "full":
-            coord_cube = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1)
-            values = [b_cube[:, :, :1].reshape((-1, 3)), b_cube[:, :, -1:].reshape((-1, 3)),
-                      b_cube[:, :1, :].reshape((-1, 3)), b_cube[:, -1:, :].reshape((-1, 3)),
-                      b_cube[:1, :, :].reshape((-1, 3)), b_cube[-1:, :, :].reshape((-1, 3)), ]
-            coords = [coord_cube[:, :, :1].reshape((-1, 3)), coord_cube[:, :, -1:].reshape((-1, 3)),
-                      coord_cube[:, :1, :].reshape((-1, 3)), coord_cube[:, -1:, :].reshape((-1, 3)),
-                      coord_cube[:1, :, :].reshape((-1, 3)), coord_cube[-1:, :, :].reshape((-1, 3)), ]
-            #
-            coords = np.concatenate(coords).astype(np.float32)
-            values = np.concatenate(values).astype(np.float32)
-        elif boundary['type'] == "potential":
-            b_cube = b_cube[:, :, 0]
-            coords, values, err = prep_b_data(b_cube, np.zeros_like(b_cube), height,
-                                              potential_boundary=True,
-                                              potential_strides=boundary['strides'])
-        elif boundary['type'] == "open":
-            # load dataset
-            coords = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1).astype(np.float32)
-            height_slices = boundary['height_slices'] if 'height_slices' in boundary else 0
-            b_cube = b_cube[:, :, height_slices]
-            coords = coords[:, :, height_slices]
-            # flatten data
-            coords = coords.reshape((-1, 3)).astype(np.float32)
-            values = b_cube.reshape((-1, 3)).astype(np.float32)
-        elif boundary['type'] == "tau":
-
-            # load dataset
-            coords = np.stack(np.mgrid[:b_cube.shape[0], :b_cube.shape[1], :b_cube.shape[2]], -1).astype(np.float32)
-            ranges = np.zeros((*coords.shape[:3], 2), dtype=np.float32)
-
-            height_mapping = {"z":  [h / 2 for h in tau_surfaces],
-                                   "z_min": [0] * len(tau_surfaces),
-                                   "z_max": [h for h in tau_surfaces]}
-            self.height_mapping = height_mapping
-            for i, (z, z_min, z_max) in enumerate(zip(height_mapping["z"], height_mapping["z_min"], height_mapping["z_max"])):
-                ranges[:, :, i, 0] = z_min
-                ranges[:, :, i, 1] = z_max
-                coords[:, :, i, 2] = z
-            #
-            if boundary['use_LOS']:
-                b_cube[:, :, 1:, 0] = None
-                b_cube[:, :, 1:, 1] = None
-            # flatten data
-            coords = coords.reshape((-1, 3)).astype(np.float32)
-            values = b_cube.reshape((-1, 3)).astype(np.float32)
-            ranges = ranges.reshape((-1, 2)).astype(np.float32)
-        else:
-            raise Exception(f'Invalid boundary condition: {boundary}. Available options: ["full", "potential", "open"]')
-
-
-        # normalize data
-        values = values / b_norm
-        # apply spatial normalization
-        coords = coords / spatial_norm
-
-        cube_shape = [*b_cube.shape[:2], height]
-        self.cube_shape = cube_shape
-
-        # prep dataset
-        # shuffle data
-        r = np.random.permutation(coords.shape[0])
-        coords = coords[r]
-        values = values[r]
-        # store data to disk
-        coords_npy_path = os.path.join(work_directory, 'coords.npy')
-        np.save(coords_npy_path, coords)
-        values_npy_path = os.path.join(work_directory, 'values.npy')
-        np.save(values_npy_path, values)
-        # create data loaders
-        batches_path = {'coords': coords_npy_path, 'values': values_npy_path}
-        if boundary['type'] == "tau":  # add ranges
-            ranges /= spatial_norm
-            ranges = ranges[r]
-            ranges_npy_path = os.path.join(work_directory, 'ranges.npy')
-            np.save(ranges_npy_path, ranges)
-            batches_path['height_ranges'] = ranges_npy_path
-
-        boundary_batch_size = int(batch_size['boundary']) if isinstance(batch_size, dict) else int(batch_size)
-        random_batch_size = int(batch_size['random']) if isinstance(batch_size, dict) else int(batch_size)
-
-        self.dataset = BatchesDataset(batches_path, boundary_batch_size)
-        self.random_dataset = RandomCoordinateDataset(cube_shape, spatial_norm, random_batch_size)
-        self.cube_dataset = CubeDataset(cube_shape, spatial_norm, batch_size=boundary_batch_size)
-        self.batches_path = batches_path
-
-    def clear(self):
-        [os.remove(f) for f in self.batches_path.values()]
-
-    def train_dataloader(self):
-        data_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                 sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
-        random_loader = DataLoader(self.random_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                   sampler=RandomSampler(self.dataset, replacement=True, num_samples=self.iterations))
-        return {'boundary': data_loader, 'random': random_loader}
-
-    def val_dataloader(self):
-        cube_loader = DataLoader(self.cube_dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                 shuffle=False)
-        boundary_loader = DataLoader(self.dataset, batch_size=None, num_workers=self.num_workers, pin_memory=True,
-                                     shuffle=False)
-        return [cube_loader, boundary_loader]
-
-
-
-
-def vector_spherical_to_cartesian(v, c):
-    vp, vt, vr = v[:, 0], v[:, 1], v[:, 2]
-    p, t, r = c[:, 0], c[:, 1], c[:, 2]
-    sin = np.sin
-    cos = np.cos
-    #
-    vx = vr * sin(t) * cos(p) + vt * cos(t) * cos(p) - vp * sin(p)
-    vy = vr * sin(t) * sin(p) + vt * cos(t) * sin(p) + vp * cos(p)
-    vz = vr * cos(t) - vt * sin(t)
-    #
-    return np.stack([vx, vy, vz], -1)
-
-
-def to_cartesian(v):
-    sin = np.sin
-    cos = np.cos
-    p, t, r = v[..., 0], v[..., 1], v[..., 2]
-    x = r * sin(t) * cos(p)
-    y = r * sin(t) * sin(p)
-    z = r * cos(t)
-    return np.stack([x, y, z], -1)
diff --git a/nf2/train/extrapolate_analytic.py b/nf2/train/extrapolate_analytic.py
index 34a03b5..76b7630 100644
--- a/nf2/train/extrapolate_analytic.py
+++ b/nf2/train/extrapolate_analytic.py
@@ -8,7 +8,7 @@
 from pytorch_lightning.loggers import WandbLogger
 
 from nf2.module import NF2Module, save
-from nf2.train.data_loader import AnalyticDataModule
+from nf2.loader.analytical import AnalyticDataModule
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--config', type=str, required=True,
diff --git a/nf2/train/loss.py b/nf2/train/loss.py
new file mode 100644
index 0000000..5151c2a
--- /dev/null
+++ b/nf2/train/loss.py
@@ -0,0 +1,275 @@
+import torch
+from torch import nn
+
+from nf2.data.util import cartesian_to_spherical, img_to_los_trv_azi, los_trv_azi_to_img
+from astropy import units as u
+
+class BaseLoss(nn.Module):
+
+    def __init__(self, name, ds_id, **kwargs):
+        super().__init__()
+        self.name = name
+        self.ds_id = ds_id
+
+
+class ForceFreeLoss(BaseLoss):
+
+    def forward(self, b, jac_matrix, *args, **kwargs):
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dx = jac_matrix[:, 1, 0]
+        dBz_dx = jac_matrix[:, 2, 0]
+        dBx_dy = jac_matrix[:, 0, 1]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dy = jac_matrix[:, 2, 1]
+        dBx_dz = jac_matrix[:, 0, 2]
+        dBy_dz = jac_matrix[:, 1, 2]
+        dBz_dz = jac_matrix[:, 2, 2]
+        #
+        rot_x = dBz_dy - dBy_dz
+        rot_y = dBx_dz - dBz_dx
+        rot_z = dBy_dx - dBx_dy
+        #
+        j = torch.stack([rot_x, rot_y, rot_z], -1)
+        jxb = torch.cross(j, b, -1)
+        # assure that the normalization does not influence the loss
+        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7).detach()
+        force_free_loss = torch.sum(jxb ** 2, dim=-1) / normalization
+        #
+        return force_free_loss.mean()
+
+
+class MagnetoStaticLoss(BaseLoss):
+
+    def forward(self, b, jac_matrix, *args, **kwargs):
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dx = jac_matrix[:, 1, 0]
+        dBz_dx = jac_matrix[:, 2, 0]
+        dP_dx = jac_matrix[:, 3, 0]
+        dBx_dy = jac_matrix[:, 0, 1]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dy = jac_matrix[:, 2, 1]
+        dP_dy = jac_matrix[:, 3, 1]
+        dBx_dz = jac_matrix[:, 0, 2]
+        dBy_dz = jac_matrix[:, 1, 2]
+        dBz_dz = jac_matrix[:, 2, 2]
+        dP_dz = jac_matrix[:, 3, 2]
+        #
+        rot_x = dBz_dy - dBy_dz
+        rot_y = dBx_dz - dBz_dx
+        rot_z = dBy_dx - dBx_dy
+        #
+        j = torch.stack([rot_x, rot_y, rot_z], -1)
+        jxb = torch.cross(j, b, -1)
+        grad_P = torch.stack([dP_dx, dP_dy, dP_dz], -1)
+        #
+        equation = jxb - grad_P
+        # assure that the normalization does not influence the loss
+        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7).detach()
+        loss = torch.sum(equation ** 2, dim=-1) / normalization
+        #
+        return loss.mean()
+
+
+class DivergenceLoss(BaseLoss):
+
+    def forward(self, jac_matrix, *args, **kwargs):
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dz = jac_matrix[:, 2, 2]
+
+        divergence_loss = (dBx_dx + dBy_dy + dBz_dz) ** 2
+
+        return divergence_loss.mean()
+
+
+class RadialLoss(BaseLoss):
+
+    def __init__(self, base_radius, Mm_per_ds, **kwargs):
+        super().__init__(**kwargs)
+        self.base_radius = (base_radius * u.solRad).to_value(u.Mm) / Mm_per_ds
+
+    def forward(self, b, coords, *args, **kwargs):
+        # radial regularization --> vanishing phi and theta components
+        radial_regularization = torch.norm(torch.cross(b, coords, dim=-1), dim=-1)
+
+        radius_weight = torch.sqrt(torch.sum(coords ** 2, dim=-1) + 1e-7)
+        radius_weight = torch.clip(radius_weight - self.base_radius, min=0)
+
+        normalization = torch.norm(b, dim=-1) * torch.norm(coords, dim=-1) + 1e-7
+        radial_regularization = radial_regularization / normalization
+
+        radial_regularization = (radial_regularization * radius_weight).mean()
+
+        return radial_regularization
+
+class PotentialLoss(BaseLoss):
+
+    def __init__(self, base_radius=None, Mm_per_ds=None, **kwargs):
+        super().__init__(**kwargs)
+        if base_radius is not None:
+            self.base_radius = (base_radius * u.solRad).to_value(u.Mm) / Mm_per_ds
+            self.solar_radius = (1 * u.solRad).to_value(u.Mm) / Mm_per_ds
+        else:
+            self.base_radius = None
+
+    def forward(self, jac_matrix, coords, *args, **kwargs):
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dx = jac_matrix[:, 1, 0]
+        dBz_dx = jac_matrix[:, 2, 0]
+        dBx_dy = jac_matrix[:, 0, 1]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dy = jac_matrix[:, 2, 1]
+        dBx_dz = jac_matrix[:, 0, 2]
+        dBy_dz = jac_matrix[:, 1, 2]
+        dBz_dz = jac_matrix[:, 2, 2]
+        #
+        rot_x = dBz_dy - dBy_dz
+        rot_y = dBx_dz - dBz_dx
+        rot_z = dBy_dx - dBx_dy
+        #
+        j = torch.stack([rot_x, rot_y, rot_z], -1)
+        potential_loss = torch.sum(j ** 2, dim=-1)
+
+        if self.base_radius is not None:
+            radius = coords.pow(2).sum(-1).pow(0.5) + 1e-7
+            radius_weight = torch.clip(radius - self.base_radius, min=0) / self.solar_radius # normalize to solar radius
+            potential_loss *= radius_weight ** 2
+
+        return potential_loss.mean()
+
+
+class EnergyGradientLoss(BaseLoss):
+
+    def __init__(self, base_radius, Mm_per_ds, **kwargs):
+        super().__init__(**kwargs)
+        self.base_radius = (base_radius * u.solRad).to_value(u.Mm) / Mm_per_ds
+        self.solar_radius = (1 * u.solRad).to_value(u.Mm) / Mm_per_ds
+        # self.asinh_stretch = nn.Parameter(torch.tensor(np.arcsinh(1e3), dtype=torch.float32), requires_grad=False)
+
+    def forward(self, b, jac_matrix, coords, *args, **kwargs):
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dx = jac_matrix[:, 1, 0]
+        dBz_dx = jac_matrix[:, 2, 0]
+        dBx_dy = jac_matrix[:, 0, 1]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dy = jac_matrix[:, 2, 1]
+        dBx_dz = jac_matrix[:, 0, 2]
+        dBy_dz = jac_matrix[:, 1, 2]
+        dBz_dz = jac_matrix[:, 2, 2]
+        # E = b^2 = b_x^2 + b_y^2 + b_z^2
+        # dE/dx = 2 * (b_x * dBx_dx + b_y * dBy_dx + b_z * dBz_dx)
+        # dE/dy = 2 * (b_x * dBx_dy + b_y * dBy_dy + b_z * dBz_dy)
+        # dE/dz = 2 * (b_x * dBx_dz + b_y * dBy_dz + b_z * dBz_dz)
+        dE_dx = 2 * (b[:, 0] * dBx_dx + b[:, 1] * dBy_dx + b[:, 2] * dBz_dx)
+        dE_dy = 2 * (b[:, 0] * dBx_dy + b[:, 1] * dBy_dy + b[:, 2] * dBz_dy)
+        dE_dz = 2 * (b[:, 0] * dBx_dz + b[:, 1] * dBy_dz + b[:, 2] * dBz_dz)
+
+        coords_spherical = cartesian_to_spherical(coords, f=torch)
+        t = coords_spherical[:, 1]
+        p = coords_spherical[:, 2]
+        dE_dr = (torch.sin(t) * torch.cos(p)) * dE_dx + \
+                (torch.sin(t) * torch.sin(p)) * dE_dy + \
+                torch.cos(p) * dE_dz
+
+        radius_weight = coords.pow(2).sum(-1).pow(0.5)
+        radius_weight = torch.clip(radius_weight - self.base_radius, min=0) / self.solar_radius # normalize to solar radius
+
+        sampled_dE_dr = dE_dr
+        energy_gradient_regularization = torch.relu(sampled_dE_dr) * radius_weight ** 2
+        energy_gradient_regularization = energy_gradient_regularization.mean()
+
+        return energy_gradient_regularization
+
+class EnergyLoss(BaseLoss):
+
+    def __init__(self, base_radius, Mm_per_ds, **kwargs):
+        super().__init__(**kwargs)
+        self.base_radius = (base_radius * u.solRad).to_value(u.Mm) / Mm_per_ds
+
+    def forward(self, b, coords, *args, **kwargs):
+        energy_loss = torch.norm(b, dim=-1).pow(2).mean()
+
+        radius_weight = coords.pow(2).sum(-1).pow(0.5)
+        radius_weight = torch.clip(radius_weight - self.base_radius, min=0)
+
+        energy_regularization = energy_loss * radius_weight ** 2
+        energy_regularization = energy_regularization.mean()
+
+        return energy_regularization
+
+class NaNLoss(BaseLoss):
+
+    def forward(self, boundary_b, b_true):
+        if torch.isnan(b_true).sum() == 0:
+            min_energy_NaNs_regularization = torch.zeros((1,), device=b_true.device)
+        else:
+            min_energy_NaNs_regularization = boundary_b[torch.isnan(b_true)].pow(2).sum() / torch.isnan(b_true).sum()
+        return min_energy_NaNs_regularization
+
+
+class LosTrvAziBoundaryLoss(BaseLoss):
+
+    def __init__(self, disambiguate=True, **kwargs):
+        super().__init__(**kwargs)
+        self.disambiguate = disambiguate
+
+    def forward(self, b, b_true, transform=None, *args, **kwargs):
+        # apply transforms
+        b_pred = torch.einsum('ijk,ik->ij', transform, b) if transform is not None else b
+        b_pred = img_to_los_trv_azi(b_pred, f=torch)
+
+        bxyz_true = los_trv_azi_to_img(b_true, ambiguous=self.disambiguate, f=torch)
+        bxyz_pred = los_trv_azi_to_img(b_pred, ambiguous=self.disambiguate, f=torch)
+
+        # compute diff
+        b_diff = bxyz_pred - bxyz_true
+        b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1))
+
+        return b_diff
+
+
+class BoundaryLoss(BaseLoss):
+
+    def forward(self, b, b_true, transform=None, b_err=None, *args, **kwargs):
+        # apply transforms
+        transformed_b = torch.einsum('ijk,ik->ij', transform, b) if transform is not None else b
+        # compute diff
+        b_err = b_err if b_err is not None else torch.zeros_like(b_true)
+        b_diff = torch.clip(torch.abs(transformed_b - b_true) - b_err, 0)
+        b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1))
+
+        assert not torch.isnan(b_diff), 'b_diff is nan'
+        return b_diff
+
+
+class HeightLoss(BaseLoss):
+
+    def forward(self, coords, original_coords, height_range, *args, **kwargs):
+        height_diff = torch.abs(coords[:, 2] - original_coords[:, 2])
+        normalization = (height_range[:, 1] - height_range[:, 0]) + 1e-6
+        height_regularization = torch.true_divide(height_diff, normalization).mean()
+
+        return height_regularization
+
+
+class MinHeightLoss(BaseLoss):
+
+    def forward(self, b_true, coords, *args, **kwargs):
+        min_height_regularization = torch.abs(coords[:, 2])
+        # min_height_regularization = min_height_regularization / (torch.norm(b_true, dim=-1) + 1e-7)
+        min_height_regularization = min_height_regularization.mean()
+        return min_height_regularization
+
+
+class FluxPreservationLoss(BaseLoss):
+
+    def forward(self, dflux_dr, *args, **kwargs):
+        flux_preservation_loss = dflux_dr.pow(2).mean()
+        return flux_preservation_loss
+
+
+class SphericalTransform(nn.Module):
+
+    def forward(self, b, b_true, transform=None):
+        b = torch.einsum('ijk,ik->ij', transform, b) if transform is not None else b
+        return b, b_true
diff --git a/nf2/train/mapping.py b/nf2/train/mapping.py
new file mode 100644
index 0000000..f6936e3
--- /dev/null
+++ b/nf2/train/mapping.py
@@ -0,0 +1,31 @@
+from nf2.data.dataset import SphereSlicesDataset, SlicesDataset, SphereDataset, CubeDataset
+from nf2.loader.base import MapDataset
+from nf2.loader.fits import FITSDataset
+from nf2.loader.spherical import SphericalSliceDataset
+from nf2.train.callback import SphericalSlicesCallback, SlicesCallback, MetricsCallback, BoundaryCallback, \
+    LosTrvAziBoundaryCallback
+
+
+def load_callbacks(data_module):
+    callbacks = []
+    Mm_per_ds = data_module.config['Mm_per_ds']
+    G_per_dB = data_module.config['G_per_dB']
+    for validation_dataset_key in data_module.validation_dataset_mapping.values():
+        ds = data_module.validation_datasets[validation_dataset_key]
+        if isinstance(ds, SphereSlicesDataset):
+            callback = SphericalSlicesCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+        elif isinstance(ds, SlicesDataset):
+            callback = SlicesCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+        elif isinstance(ds, SphereDataset) or isinstance(ds, CubeDataset):
+            callback = MetricsCallback(validation_dataset_key, G_per_dB, Mm_per_ds)
+        elif isinstance(ds, SphericalSliceDataset) or isinstance(ds, FITSDataset):
+            callback = BoundaryCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+        elif isinstance(ds, MapDataset):
+            if ds.los_trv_azi:
+                callback = LosTrvAziBoundaryCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+            else:
+                callback = BoundaryCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+        else:
+            raise NotImplementedError(f'Data set {type(ds)} not implemented for validation.')
+        callbacks += [callback]
+    return callbacks
diff --git a/nf2/train/model.py b/nf2/train/model.py
index edd87c6..31348cb 100644
--- a/nf2/train/model.py
+++ b/nf2/train/model.py
@@ -1,6 +1,32 @@
+import numpy as np
 import torch
+from astropy import units as u
+from matplotlib import pyplot as plt
+from sunpy.coordinates import frames
+from sunpy.map import Map, all_coordinates_from_map
 from torch import nn
-from torch.nn import Embedding
+
+from nf2.data.util import cartesian_to_spherical, spherical_to_cartesian
+
+
+class Swish(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.beta = nn.Parameter(torch.tensor(1., dtype=torch.float32), requires_grad=True)
+
+    def forward(self, x):
+        return x * torch.sigmoid(self.beta * x)
+
+
+class Sinsh(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.beta = nn.Parameter(torch.randn((1,), dtype=torch.float32), requires_grad=True)
+
+    def forward(self, x):
+        return x * torch.sin(self.beta * x)
 
 
 class Sine(nn.Module):
@@ -11,76 +37,143 @@ def __init__(self, w0=1.):
     def forward(self, x):
         return torch.sin(self.w0 * x)
 
-class HeightMappingModel(nn.Module):
 
-    def __init__(self, in_coords, dim, positional_encoding=True):
+class HeightTransformModel(nn.Module):
+
+    def __init__(self, in_coords, ds_id, validation_ds_id=[], dim=256, positional_encoding=True, ):
         super().__init__()
         if positional_encoding:
-            posenc = PositionalEncoding(8, 20)
+            self.posenc = PositionalEncoding(20, in_coords)
             d_in = nn.Linear(in_coords * 40, dim)
-            self.d_in = nn.Sequential(posenc, d_in)
+            self.d_in = nn.Sequential(self.posenc, d_in)
         else:
             self.d_in = nn.Linear(in_coords, dim)
         lin = [nn.Linear(dim, dim) for _ in range(4)]
         self.linear_layers = nn.ModuleList(lin)
         self.d_out = nn.Linear(dim, 1)
         self.activation = Sine()
+        self.ds_id = ds_id if isinstance(ds_id, list) else [ds_id]
+        self.validation_ds_id = validation_ds_id if isinstance(validation_ds_id, list) else [validation_ds_id]
+
+    def forward(self, batch):
+        coords = torch.cat([batch[ds_id]['coords'] for ds_id in self.ds_id], 0)
+        height_range = torch.cat([batch[ds_id]['height_range'] for ds_id in self.ds_id], 0)
+        transformed_coords = self.transform_coords(coords, height_range)
+
+        coord_idx = 0
+        for ds_id in self.ds_id:
+            n_coords = batch[ds_id]['coords'].shape[0]
+            batch[ds_id]['original_coords'] = batch[ds_id]['coords']
+            batch[ds_id]['coords'] = transformed_coords[coord_idx:coord_idx + n_coords]
+            coord_idx += n_coords
+        return batch
+
+    def transform_batch(self, batch):
+        coords = batch['coords']
+        transformed_coords = self.transform_coords(**batch)
+        batch['coords'] = transformed_coords
+        batch['original_coords'] = coords
+        return batch
 
-    def forward(self, x, height_range):
-        input_coords = x
-        x = self.activation(self.d_in(x))
+    def transform_coords(self, coords, height_range, **kwargs):
+        x = self.activation(self.d_in(coords))
         for l in self.linear_layers:
             x = self.activation(l(x))
         z_coords = torch.sigmoid(self.d_out(x)) * (height_range[:, 1:2] - height_range[:, 0:1]) + height_range[:, 0:1]
-        # shifted_z_coords = input_coords[:, 2:3] * (1 + z_shift) # max shift in dependence of height estimate
-        output_coords = torch.cat([input_coords[:, :2], z_coords], -1)
+        output_coords = torch.cat([coords[:, :2], z_coords], -1)
         return output_coords
 
-class BModel(nn.Module):
 
-    def __init__(self, in_coords, out_values, dim, pos_encoding=False):
+class RadialTransformModel(nn.Module):
+
+    def __init__(self, in_coords, dim, positional_encoding=True, ds_ids=[]):
         super().__init__()
-        if pos_encoding:
-            posenc = PositionalEncoding(8, 20)
+        if positional_encoding:
+            posenc = PositionalEncoding(20, in_coords)
             d_in = nn.Linear(in_coords * 40, dim)
             self.d_in = nn.Sequential(posenc, d_in)
         else:
             self.d_in = nn.Linear(in_coords, dim)
-        lin = [nn.Linear(dim, dim) for _ in range(8)]
+        lin = [nn.Linear(dim, dim) for _ in range(4)]
         self.linear_layers = nn.ModuleList(lin)
-        self.d_out = nn.Linear(dim, out_values)
-        self.activation = Sine()  # torch.tanh
+        self.d_out = nn.Linear(dim, 1)
+        self.activation = Sine()
+        self.ds_ids = ds_ids
+        self.observer_transformer = ObserverTransformer()
 
-    def forward(self, x):
-        x = self.activation(self.d_in(x))
+    def forward(self, batch):
+        for ds_id in self.ds_ids:
+            transformed_coords = self.transform(batch[ds_id]['coords'], batch[ds_id]['obs_coords'],
+                                                batch[ds_id]['height_range'])
+            batch[ds_id]['coords'] = transformed_coords
+            batch[ds_id]['original_coords'] = coords
+
+    def transform(self, coords, obs_coords, height_range, **kwargs):
+        coords = self.observer_transformer.transform(coords, obs_coords)
+
+        x = self.activation(self.d_in(coords))
         for l in self.linear_layers:
             x = self.activation(l(x))
-        x = self.d_out(x)
-        return x
+        z_coords = torch.sigmoid(self.d_out(x)) * (height_range[:, 1:2] - height_range[:, 0:1]) + height_range[:, 0:1]
+        output_coords = torch.cat([coords[:, :2], z_coords], -1)
 
-class VectorPotentialModel(nn.Module):
+        output_coords = self.observer_transformer.inverse_transform(output_coords, obs_coords)
 
-    def __init__(self, in_coords, dim, pos_encoding=False):
+        return output_coords
+
+
+class GenericModel(nn.Module):
+
+    def __init__(self, in_coords, out_coords, dim=256, encoding=None, activation='swish'):
         super().__init__()
-        if pos_encoding:
-            posenc = PositionalEncoding(8, 20)
+        if encoding is None or encoding == 'none':
+            self.d_in = nn.Linear(in_coords, dim)
+        elif encoding == 'positional':
+            posenc = PositionalEncoding(20, in_coords)
             d_in = nn.Linear(in_coords * 40, dim)
             self.d_in = nn.Sequential(posenc, d_in)
         else:
-            self.d_in = nn.Linear(in_coords, dim)
+            raise NotImplementedError(f'Unknown encoding {encoding}')
         lin = [nn.Linear(dim, dim) for _ in range(8)]
         self.linear_layers = nn.ModuleList(lin)
-        self.d_out = nn.Linear(dim, 3)
-        self.activation = Sine()  # torch.tanh
+        self.d_out = nn.Linear(dim, out_coords)
+        activation_mapping = {'relu': nn.ReLU, 'swish': Swish, 'tanh': nn.Tanh, 'sine': Sine}
+        activation_f = activation_mapping[activation]
+        self.in_activation = activation_f()
+        self.activations = nn.ModuleList([activation_f() for _ in range(8)])
 
     def forward(self, x):
-        coord = x
-        x = self.activation(self.d_in(x))
-        for l in self.linear_layers:
-            x = self.activation(l(x))
-        a = self.d_out(x)
+        radius = torch.norm(x, dim=-1, keepdim=True)
+        x = self.in_activation(self.d_in(x))
+        for l, a in zip(self.linear_layers, self.activations):
+            x = a(l(x))
+        x = self.d_out(x)
+        return x
+
+
+class BModel(GenericModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(3, 3, **kwargs)
+
+    def forward(self, coords, compute_jacobian=True):
+        b = super().forward(coords)
+        out_dict = {'b': b}
+        if compute_jacobian:
+            jac_matrix = jacobian(b, coords)
+            out_dict['jac_matrix'] = jac_matrix
+        return out_dict
+
+
+class VectorPotentialModel(GenericModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(3, 3, **kwargs)
+
+    def forward(self, coords, compute_jacobian=True):
+        a = super().forward(coords)
         #
-        jac_matrix = jacobian(a, coord)
+        jac_matrix = jacobian(a, coords)
         dAy_dx = jac_matrix[:, 1, 0]
         dAz_dx = jac_matrix[:, 2, 0]
         dAx_dy = jac_matrix[:, 0, 1]
@@ -91,39 +184,143 @@ def forward(self, x):
         rot_y = dAx_dz - dAz_dx
         rot_z = dAy_dx - dAx_dy
         b = torch.stack([rot_x, rot_y, rot_z], -1)
+        out_dict = {'b': b, 'a': a}
         #
-        return b
+        if compute_jacobian:
+            jac_matrix = jacobian(b, coords)
+            out_dict['jac_matrix'] = jac_matrix
+        #
+        return out_dict
+
+
+class PressureModel(GenericModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(3, 1, **kwargs)
+        self.softplus = nn.Softplus()
+
+    def forward(self, x):
+        p = super().forward(x)
+        p = 10 ** p
+        return {'p': p}
+
+
+class MagnetoStaticModel(GenericModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(3, 4)
+
+    def forward(self, coords, compute_jacobian=True):
+        model_out = super().forward(coords)
+        b = model_out[:, :3]
+        p = 10 ** model_out[:, 3:]
+        out_dict = {'b': b, 'p': p}
+        if compute_jacobian:
+            jac_matrix = jacobian(model_out, coords)
+            out_dict['jac_matrix'] = jac_matrix
+        return out_dict
+
+
+class MagnetoStaticModelV2(nn.Module):
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.b_model = VectorPotentialModel(**kwargs)
+        self.p_model = PressureModel(**kwargs)
+
+    def forward(self, coords, compute_jacobian=True):
+        b_dict = self.b_model(coords)
+        p_dict = self.p_model(coords)
+        out_dict = {**b_dict, **p_dict}
+        if compute_jacobian:
+            b = b_dict['b']
+            jac_b_matrix = jacobian(b, coords)
+            p = p_dict['p']
+            jac_p_matrix = jacobian(p, coords)
+            jac_matrix = torch.cat([jac_b_matrix, jac_p_matrix], -2)
+            out_dict['jac_matrix'] = jac_matrix
+
+        return out_dict
+
 
 class PositionalEncoding(nn.Module):
-    """
-    Positional Encoding of the input coordinates.
-
-    encodes x to (..., sin(2^k x), cos(2^k x), ...)
-    k takes "num_freqs" number of values equally spaced between [0, max_freq]
-    """
-
-    def __init__(self, max_freq, num_freqs):
-        """
-        Args:
-            max_freq (int): maximum frequency in the positional encoding.
-            num_freqs (int): number of frequencies between [0, max_freq]
-        """
+
+    def __init__(self, num_freqs, in_features):
         super().__init__()
-        freqs = 2 ** torch.linspace(0, max_freq, num_freqs)
-        self.register_buffer("freqs", freqs)  # (num_freqs)
+        frequencies = torch.randn(num_freqs, in_features)
+        self.frequencies = nn.Parameter(frequencies[None], requires_grad=True)
 
     def forward(self, x):
-        """
-        Inputs:
-            x: (batch, num_samples, in_features)
-        Outputs:
-            out: (batch, num_samples, 2*num_freqs*in_features)
-        """
-        x_proj = x.unsqueeze(dim=-2) * self.freqs.unsqueeze(dim=-1)  # (num_rays, num_samples, num_freqs, in_features)
-        x_proj = x_proj.reshape(*x.shape[:-1], -1)  # (num_rays, num_samples, num_freqs*in_features)
-        out = torch.cat([torch.sin(x_proj), torch.cos(x_proj)],
-                        dim=-1)  # (num_rays, num_samples, 2*num_freqs*in_features)
-        return out
+        encoded = x[:, None, :] * torch.pi * 2 ** self.frequencies
+        encoded = encoded.reshape(x.shape[0], -1)
+        encoded = torch.cat([torch.sin(encoded), torch.cos(encoded)], -1)
+        return encoded
+
+
+class ObserverTransformer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, coords, obs_coord):
+        return self.transform(coords, obs_coord)
+
+    def transform(self, coords, obs_coord):
+        # transform coords to observer frame
+        coords = coords - obs_coord
+        coords = cartesian_to_spherical(coords, f=torch)
+        return coords
+
+    def inverse_transform(self, coords, obs_coord):
+        # transform coords to solar frame
+        coords = spherical_to_cartesian(coords, f=torch)
+        coords = coords + obs_coord
+        return coords
+
+
+def image_to_spherical_matrix(lon, lat, latc, lonc, pAng, sin=np.sin, cos=np.cos):
+    a11 = -sin(latc) * sin(pAng) * sin(lon - lonc) + cos(pAng) * cos(lon - lonc)
+    a12 = sin(latc) * cos(pAng) * sin(lon - lonc) + sin(pAng) * cos(lon - lonc)
+    a13 = -cos(latc) * sin(lon - lonc)
+    a21 = -sin(lat) * (sin(latc) * sin(pAng) * cos(lon - lonc) + cos(pAng) * sin(lon - lonc)) - cos(lat) * cos(
+        latc) * sin(pAng)
+    a22 = sin(lat) * (sin(latc) * cos(pAng) * cos(lon - lonc) - sin(pAng) * sin(lon - lonc)) + cos(lat) * cos(
+        latc) * cos(pAng)
+    a23 = -cos(latc) * sin(lat) * cos(lon - lonc) + sin(latc) * cos(lat)
+    a31 = cos(lat) * (sin(latc) * sin(pAng) * cos(lon - lonc) + cos(pAng) * sin(lon - lonc)) - sin(lat) * cos(
+        latc) * sin(pAng)
+    a32 = -cos(lat) * (sin(latc) * cos(pAng) * cos(lon - lonc) - sin(pAng) * sin(lon - lonc)) + sin(lat) * cos(
+        latc) * cos(pAng)
+    a33 = cos(lat) * cos(latc) * cos(lon - lonc) + sin(lat) * sin(latc)
+
+    # a_matrix = np.stack([a11, a12, a13, a21, a22, a23, a31, a32, a33], axis=-1)
+    a_matrix = np.stack([a31, a32, a33, a21, a22, a23, a11, a12, a13], axis=-1)
+    a_matrix = a_matrix.reshape((*a_matrix.shape[:-1], 3, 3))
+    return a_matrix
+
+
+def calculate_current(b, coords, jac_matrix=None):
+    jac_matrix = jacobian(b, coords) if jac_matrix is None else jac_matrix
+    j = calculate_current_from_jacobian(jac_matrix)
+    return j
+
+def calculate_current_from_jacobian(jac_matrix, f=torch):
+    dBx_dx = jac_matrix[..., 0, 0]
+    dBy_dx = jac_matrix[..., 1, 0]
+    dBz_dx = jac_matrix[..., 2, 0]
+    dBx_dy = jac_matrix[..., 0, 1]
+    dBy_dy = jac_matrix[..., 1, 1]
+    dBz_dy = jac_matrix[..., 2, 1]
+    dBx_dz = jac_matrix[..., 0, 2]
+    dBy_dz = jac_matrix[..., 1, 2]
+    dBz_dz = jac_matrix[..., 2, 2]
+    #
+    rot_x = dBz_dy - dBy_dz
+    rot_y = dBx_dz - dBz_dx
+    rot_z = dBy_dx - dBx_dy
+    #
+    j = f.stack([rot_x, rot_y, rot_z], -1)
+    return j
 
 
 def jacobian(output, coords):
@@ -132,4 +329,148 @@ def jacobian(output, coords):
                                       retain_graph=True, create_graph=True, allow_unused=True)[0]
                   for i in range(output.shape[1])]
     jac_matrix = torch.stack(jac_matrix, dim=1)
-    return jac_matrix
\ No newline at end of file
+    return jac_matrix
+
+
+if __name__ == '__main__':
+    B_field_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.field.fits')
+    B_az_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.azimuth.fits')
+    B_disamb_map = Map(
+        '/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/hmi.B_720s.20140902_060000_TAI.disambig.fits')
+    B_in_map = Map(
+        '/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/hmi.b_720s.20140902_060000_TAI.inclination.fits')
+
+    fld = B_field_map.data
+    inc = np.deg2rad(B_in_map.data)
+    azi = np.deg2rad(B_az_map.data)
+    amb = B_disamb_map.data
+    # disambiguate
+    amb_weak = 2
+    condition = (amb.astype((int)) >> amb_weak).astype(bool)
+    azi[condition] += np.pi
+
+    sin = np.sin
+    cos = np.cos
+    B_x = - fld * sin(inc) * sin(azi)
+    B_y = fld * sin(inc) * cos(azi)
+    B_z = fld * cos(inc)
+
+    fig, axs = plt.subplots(3, 2, figsize=(6, 9))
+    axs[0, 0].imshow(B_x, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[0, 0].set_title('Bx')
+    axs[0, 1].imshow(B_field_map.data, cmap='gray', origin='lower', vmin=0, vmax=1000)
+    axs[0, 1].set_title('B')
+    #
+    axs[1, 0].imshow(B_y, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[1, 0].set_title('By')
+    axs[1, 1].imshow(B_in_map.data, cmap='gray', origin='lower', vmin=0, vmax=180)
+    axs[1, 1].set_title('Inclination')
+    #
+    axs[2, 0].imshow(B_z, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[2, 0].set_title('Bz')
+    axs[2, 1].imshow(B_az_map.data, cmap='gray', origin='lower', vmin=0, vmax=360)
+    axs[2, 1].set_title('Azimuth')
+    plt.tight_layout()
+    plt.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/image_frame.jpg', dpi=150)
+    plt.close()
+
+    B = np.stack([B_x, B_y, B_z], -1)
+    latc, lonc = np.deg2rad(B_field_map.meta['CRLT_OBS']), np.deg2rad(B_field_map.meta['CRLN_OBS'])
+
+    coords = all_coordinates_from_map(B_field_map).transform_to(frames.HeliographicCarrington)
+    lat, lon = coords.lat.to(u.rad).value, coords.lon.to(u.rad).value
+
+    pAng = -np.deg2rad(B_field_map.meta['CROTA2'])
+
+    a_matrix = image_to_spherical_matrix(lon, lat, latc, lonc, pAng=pAng)
+    Brtp = np.einsum('...ij,...j->...i', a_matrix, B)
+    bp = Brtp[..., 2]
+    bt = -Brtp[..., 1]
+    br = Brtp[..., 0]
+
+    # to grid coords
+
+    ref_B_p = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bp.fits').data
+    ref_B_t = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Bt.fits').data
+    ref_B_r = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.Br.fits').data
+
+    # plot
+    fig, axs = plt.subplots(3, 2, figsize=(6, 9))
+    axs[0, 0].imshow(bp, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[0, 0].set_title('Bp')
+    axs[0, 1].imshow(ref_B_p, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[0, 1].set_title('ref Bp')
+    axs[1, 0].imshow(bt, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[1, 0].set_title('Bt')
+    axs[1, 1].imshow(ref_B_t, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[1, 1].set_title('ref Bt')
+    axs[2, 0].imshow(br, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[2, 0].set_title('Br')
+    axs[2, 1].imshow(ref_B_r, cmap='gray', origin='lower', vmin=-1000, vmax=1000)
+    axs[2, 1].set_title('ref Br')
+    plt.tight_layout()
+    plt.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/ptr_frame.jpg', dpi=150)
+    plt.close()
+
+    lat_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.lat.fits')
+    lon_map = Map('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/hmi.b_720s.20140902_060000_TAI.lon.fits')
+
+    fig, axs = plt.subplots(2, 2, figsize=(6, 6))
+    axs[0, 0].imshow(np.deg2rad(lat_map.data), origin='lower', vmin=-np.pi / 2, vmax=np.pi / 2)
+    axs[0, 0].set_title('lat')
+    axs[0, 1].imshow(lat, origin='lower', vmin=-np.pi / 2, vmax=np.pi / 2)
+    axs[0, 1].set_title('lat ref')
+    axs[1, 0].imshow(np.deg2rad(lon_map.data), origin='lower', vmin=-np.pi, vmax=np.pi)
+    axs[1, 0].set_title('lon')
+    axs[1, 1].imshow(lon, origin='lower', vmin=-np.pi, vmax=np.pi)
+    axs[1, 1].set_title('lon ref')
+    plt.tight_layout()
+    plt.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/latlon_frame.jpg', dpi=150)
+    plt.close()
+
+    # transform back to az and inc
+    ai_matrix = np.linalg.inv(a_matrix)
+    B = np.einsum('...ij,...j->...i', ai_matrix, Brtp)
+    # B_x = - fld * sin(inc) * sin(azi)
+    # B_y = fld * sin(inc) * cos(azi)
+    # B_z = fld * cos(inc)
+    eps = 1e-7
+    fld_rec = np.sqrt(B[..., 0] ** 2 + B[..., 1] ** 2 + B[..., 2] ** 2)
+    inc_rec = np.arccos(np.clip(B[..., 2] / (fld + eps), a_min=-1 + eps, a_max=1 - eps))
+    azi_rec = np.arctan2(B[..., 0], -B[..., 1] + eps)
+
+    # plot
+    fig, axs = plt.subplots(3, 2, figsize=(6, 9))
+    axs[0, 0].imshow(fld_rec, cmap='gray', origin='lower', vmin=0, vmax=1000)
+    axs[0, 0].set_title('B')
+    axs[0, 1].imshow(fld, cmap='gray', origin='lower', vmin=0, vmax=1000)
+    axs[0, 1].set_title('ref B')
+    axs[1, 0].imshow(inc_rec, cmap='gray', origin='lower', vmin=0, vmax=np.pi / 2)
+    axs[1, 0].set_title('inc')
+    axs[1, 1].imshow(inc, cmap='gray', origin='lower', vmin=0, vmax=np.pi / 2)
+    axs[1, 1].set_title('ref inc')
+    axs[2, 0].imshow(azi_rec, cmap='gray', origin='lower', vmin=-np.pi, vmax=np.pi)
+    axs[2, 0].set_title('azi')
+    axs[2, 1].imshow(azi - np.pi, cmap='gray', origin='lower', vmin=-np.pi, vmax=np.pi)
+    axs[2, 1].set_title('ref azi')
+    plt.tight_layout()
+    plt.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/azinc_frame.jpg', dpi=150)
+    plt.close()
+
+    # plot
+    fig, axs = plt.subplots(3, 2, figsize=(6, 9))
+    axs[0, 0].imshow(fld_rec, cmap='gray', origin='lower', vmin=0, vmax=1000)
+    axs[0, 0].set_title('B')
+    axs[0, 1].imshow(fld, cmap='gray', origin='lower', vmin=0, vmax=1000)
+    axs[0, 1].set_title('ref B')
+    axs[1, 0].imshow(inc_rec, cmap='gray', origin='lower', vmin=0, vmax=np.pi / 2)
+    axs[1, 0].set_title('inc')
+    axs[1, 1].imshow(inc, cmap='gray', origin='lower', vmin=0, vmax=np.pi / 2)
+    axs[1, 1].set_title('ref inc')
+    axs[2, 0].imshow(np.cos(2 * azi_rec), cmap='gray', origin='lower', vmin=-1, vmax=1)
+    axs[2, 0].set_title('azi')
+    axs[2, 1].imshow(np.cos(2 * np.deg2rad(B_az_map.data)), cmap='gray', origin='lower', vmin=-1, vmax=1)
+    axs[2, 1].set_title('ref azi')
+    plt.tight_layout()
+    fig.savefig('/gpfs/gpfs0/robert.jarolim/data/nf2/fd_2154/full_disk/disamb.jpg', dpi=150)
+    plt.close()
diff --git a/nf2/train/module.py b/nf2/train/module.py
index 8545690..93ed642 100644
--- a/nf2/train/module.py
+++ b/nf2/train/module.py
@@ -1,45 +1,57 @@
+import copy
 import logging
 
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-import wandb
 from pytorch_lightning import LightningModule
 from torch.optim.lr_scheduler import ExponentialLR
 
-from nf2.train.model import BModel, jacobian, VectorPotentialModel, HeightMappingModel
+from nf2.train.loss import *
+from nf2.train.model import BModel, VectorPotentialModel, HeightTransformModel, MagnetoStaticModel, MagnetoStaticModelV2
 
 
 class NF2Module(LightningModule):
 
-    def __init__(self, validation_settings, dim=256, lambda_b={'start': 1e3, 'end': 1, 'iterations': 1e5},
-                 lambda_div=0.1, lambda_ff=0.1, lambda_height_reg=1e-3, lambda_min_energy_nans=1e-3, lambda_min_energy=0,
-                 lr_params = {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}, meta_path=None,
-                 positional_encoding=False, use_vector_potential=False, use_height_mapping=False, **kwargs):
-        """Magnetic field extrapolations trainer
-
-        :param dim: number of neurons per layer (8 layers).
-        :param positional_encoding: use positional encoding.
-        :param use_vector_potential: derive the magnetic field from a vector potential.
-        :param lambda_div: weighting parameter for divergence freeness of the simulation.
-        :param lambda_ff: weighting parameter for force freeness of the simulation.
-        :param decay_iterations: decay weighting for boundary condition (lambda_B=1000) over n iterations to 1.
-        :param meta_path: start from a pre-learned simulation state.
-        :param work_directory: directory to store scratch data (prepared batches).
+    def __init__(self, validation_mapping, data_config, model_kwargs=None, loss_config=None,
+                 lr_params={"start": 5e-4, "end": 5e-5, "iterations": 1e5},
+                 coordinate_transform={"type": None},
+                 meta_path=None, **kwargs):
+        """
+        The main module for training the neural field model.
+
+        Args:
+            validation_mapping (dict): Dictionary mapping validation dataset indices to data loader indices.
+            dim (int): Dimension of the neural field model.
+            loss_config (list): List of dictionaries containing the loss type and lambda value.
+            lr_params (dict): Dictionary containing the start, end and iterations for the learning rate scheduler.
+            meta_path (str): Path to a meta state file.
+            use_positional_encoding (bool): Whether to use positional encoding.
+            model (str): Model type, one of ['b', 'vector_potential', 'flux'].
+            **kwargs: Additional keyword arguments.
         """
         super().__init__()
+        model_kwargs = model_kwargs if model_kwargs is not None else {'type': 'b', 'dim': 256}
         # init model
-        if use_vector_potential:
-            model = VectorPotentialModel(3, dim, pos_encoding=positional_encoding)
+        model_kwargs = copy.deepcopy(model_kwargs)
+        model_type = model_kwargs.pop('type')
+        if model_type == 'b':
+            model = BModel(**model_kwargs)
+        elif model_type == 'vector_potential':
+            model = VectorPotentialModel(**model_kwargs)
+        elif model_type == 'magneto_static':
+            model = MagnetoStaticModel(**model_kwargs)
+        elif model_type == 'magneto_static_v2':
+            model = MagnetoStaticModelV2(**model_kwargs)
         else:
-            model = BModel(3, 3, dim, pos_encoding=positional_encoding)
-        if use_height_mapping:
-            self.height_mapping_model = HeightMappingModel(3, dim)
+            raise ValueError(f"Invalid model: {model_type}, must be in ['b', 'vector_potential', 'flux']")
+
+        # init coordinate mapping model
+        coordinate_transform = copy.deepcopy(coordinate_transform)
+        transform_type = coordinate_transform.pop('type')
+        if transform_type is None or transform_type == 'none':
+            transform_module = None
+        elif transform_type == 'height':
+            transform_module = HeightTransformModel(3, **coordinate_transform)
         else:
-            self.height_mapping_model = None
-        self.model = model
-        self.validation_settings = validation_settings
-        self.lr_params = lr_params
+            raise ValueError(f"Invalid coordinate mapping model: {transform_type}")
 
         # load meta state
         if meta_path:
@@ -47,241 +59,243 @@ def __init__(self, validation_settings, dim=256, lambda_b={'start': 1e3, 'end':
                 if meta_path.endswith('nf2') else torch.load(meta_path)['m']
             model.load_state_dict(state_dict)
             logging.info('Loaded meta state: %s' % meta_path)
-        # init
-        if isinstance(lambda_b, dict):
-            self.register_buffer('lambda_B', torch.tensor(lambda_b['start'], dtype=torch.float32))
-            self.register_buffer('lambda_B_gamma', torch.tensor((lambda_b['end'] / lambda_b['start']) ** (1 / lambda_b['iterations']) \
-                if lambda_b['iterations'] > 0 else 0, dtype=torch.float32))
-            self.register_buffer('lambda_B_end', torch.tensor(lambda_b['end'], dtype=torch.float32))
-        else:
-            self.register_buffer('lambda_B', torch.tensor(lambda_b, dtype=torch.float32))
-            self.register_buffer('lambda_B_gamma', torch.tensor(1, dtype=torch.float32))
-            self.register_buffer('lambda_B_end', torch.tensor(lambda_b, dtype=torch.float32))
-        self.register_buffer('lambda_div', torch.tensor(lambda_div, dtype=torch.float32))
-        self.register_buffer('lambda_ff', torch.tensor(lambda_ff, dtype=torch.float32))
-        self.register_buffer('lambda_height_reg', torch.tensor(lambda_height_reg, dtype=torch.float32))
-        self.register_buffer('lambda_min_energy_nans', torch.tensor(lambda_min_energy_nans, dtype=torch.float32))
-        self.register_buffer('lambda_min_energy', torch.tensor(lambda_min_energy, dtype=torch.float32))
+
+        # init loss modules
+        if loss_config is None:
+            logging.info('Using default loss configuration')
+            loss_config = [
+                {"type": "boundary", "lambda": {"start": 1e3, "end": 1, "iterations": 1e5},
+                 'ds_id': ['boundary_01', 'potential']},
+                {"type": "force_free", "lambda": 1e-1},
+                {"type": "divergence", "lambda": 1e-1},
+            ]
+
+        # mapping
+        loss_module_mapping = {'boundary': BoundaryLoss, 'boundary_los_trv_azi': LosTrvAziBoundaryLoss,
+                               'divergence': DivergenceLoss, 'force_free': ForceFreeLoss, 'potential': PotentialLoss,
+                               'height': HeightLoss, 'NaNs': NaNLoss, 'radial': RadialLoss,
+                               'min_height': MinHeightLoss, 'energy_gradient': EnergyGradientLoss, 'energy': EnergyLoss,
+                               'flux_preservation': FluxPreservationLoss, 'magneto_static': MagnetoStaticLoss}
+        # init lambdas and loss modules
+        scheduled_lambdas = {}
+        lambdas = {}
+        loss_modules = {}
+        loss_config = copy.deepcopy(loss_config)
+        for config in loss_config:
+            k = config.pop('type')
+            l = config.pop('lambda')
+            # update dataset id
+            ds_id = config.pop('ds_id', 'all')
+            config['ds_id'] = ds_id
+            # update name
+            name = config.pop('name', None)
+            if name is None:
+                ds_str = '_'.join(ds_id) if isinstance(ds_id, list) else ds_id
+                name = f'{ds_str}--{k}' if ds_str != 'all' else k
+            config['name'] = name
+            if isinstance(l, dict):
+                value = torch.tensor(l['start'], dtype=torch.float32)
+                lt = 'exponential' if 'type' not in l else l['type']
+                if lt == 'exponential':
+                    gamma = torch.tensor((l['end'] / l['start']) ** (1 / l['iterations']), dtype=torch.float32)
+                elif lt == 'linear':
+                    gamma = torch.tensor((l['end'] - l['start']) / l['iterations'], dtype=torch.float32)
+                else:
+                    raise ValueError(f"Invalid lambda type: {lt}, must be in ['exponential', 'linear']")
+                end = torch.tensor(l['end'], dtype=torch.float32)
+                scheduled_lambdas[name] = {'gamma': gamma, 'end': end, 'type': lt}
+            else:
+                value = torch.tensor(l, dtype=torch.float32)
+            assert name not in lambdas, f"Duplicate name for loss: {name}"
+            lambdas[name] = value
+            # additional kwargs
+            loss_module = {name: loss_module_mapping[k](**config, **data_config)}
+            loss_modules.update(loss_module)
+
+        self.model = model
+        self.transform_module = transform_module
+        self.validation_mapping = validation_mapping
+        self.lr_params = lr_params
+        self.scheduled_lambdas = nn.ParameterDict(scheduled_lambdas)
+        self.lambdas = nn.ParameterDict(lambdas)
         #
-        self.use_vector_potential = use_vector_potential
-        self.use_height_mapping = use_height_mapping
+        self.validation_outputs = {}
+        self.loss_modules = nn.ModuleDict(loss_modules)
 
     def configure_optimizers(self):
         parameters = list(self.model.parameters())
-        if self.use_height_mapping:
-            parameters += list(self.height_mapping_model.parameters())
+        if self.transform_module is not None:
+            parameters += list(self.transform_module.parameters())
         if isinstance(self.lr_params, dict):
             lr_start = self.lr_params['start']
             lr_end = self.lr_params['end']
-            decay_iterations = self.lr_params['decay_iterations']
-        else:
+            iterations = self.lr_params['iterations']
+        elif isinstance(self.lr_params, (float, int)):
             lr_start = self.lr_params
             lr_end = self.lr_params
-            decay_iterations = 1
-        self.optimizer = torch.optim.Adam(parameters, lr=lr_start)
-        self.scheduler = ExponentialLR(self.optimizer, gamma=(lr_end / lr_start) ** (1 / decay_iterations))
+            iterations = 1
+            self.lr_params = {'start': lr_start, 'end': lr_end, 'iterations': iterations}
+        else:
+            raise ValueError(f"Invalid lr_params: {self.lr_params}, must be dict or float/int")
+        optimizer = torch.optim.Adam(parameters, lr=lr_start)
+        scheduler = ExponentialLR(optimizer, gamma=(lr_end / lr_start) ** (1 / iterations))
 
-        return [self.optimizer], [self.scheduler]
+        return [optimizer], [scheduler]
 
     def training_step(self, batch, batch_nb):
-        boundary_batch = batch['boundary']
-        boundary_coords = boundary_batch['coords']
-        b_true = boundary_batch['values']
-        b_err = boundary_batch['error'] if 'error' in boundary_batch else 0
-        boundary_coords.requires_grad = True
-        if self.use_height_mapping:
-            original_coords = boundary_coords
-            boundary_coords = self.height_mapping_model(boundary_coords, boundary_batch['height_ranges'])
-
-        random_coords = batch['random']
-        random_coords.requires_grad = True
-        # concatenate boundary and random points
-        n_boundary_coords = boundary_coords.shape[0]
-        coords = torch.cat([boundary_coords, random_coords], 0)
+        loader_keys = list(batch.keys())
+
+        # set requires grad for coords
+        for k in loader_keys:
+            batch[k]['coords'].requires_grad = True
+
+        # transform batch
+        if self.transform_module is not None:
+            batch = self.transform_module(batch)
+
+        # concatenate all points
+        coords = torch.cat([batch[k]['coords'] for k in loader_keys], 0)
 
         # forward step
-        b = self.model(coords)
-
-        # compute boundary loss
-        boundary_b = b[:n_boundary_coords]
-        b_diff = torch.clip(torch.abs(boundary_b - b_true) - b_err, 0)
-        b_diff = torch.mean(torch.nanmean(b_diff.pow(2), -1) * 3) # b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1))
-
-        # minimum energy for nan components
-        min_energy_NaNs_regularization = boundary_b[torch.isnan(b_true)].pow(2).sum() / (torch.isnan(b_true).sum() + 1e-6)
-
-        min_energy_regularization = torch.mean(torch.nansum(b[n_boundary_coords:].pow(2), -1))
-
-        # compute div and ff loss
-        divergence_loss, force_loss = calculate_loss(b, coords)
-        divergence_loss, force_loss = divergence_loss.mean(), force_loss.mean()
-        loss = b_diff * self.lambda_B + \
-               divergence_loss * self.lambda_div + \
-               force_loss * self.lambda_ff + \
-               min_energy_NaNs_regularization * self.lambda_min_energy_nans + \
-               min_energy_regularization * self.lambda_min_energy
-        loss_dict = {'b_diff': b_diff, 'divergence': divergence_loss, 'force-free': force_loss,
-                     'min_energy_NaNs_regularization': min_energy_NaNs_regularization,
-                     'min_energy_regularization': min_energy_regularization, }
-        if self.use_height_mapping:
-            height_diff = torch.abs(boundary_coords[:, 2] - original_coords[:, 2])
-            normalization = (boundary_batch['height_ranges'][:, 1] - boundary_batch['height_ranges'][:, 0]) + 1e-6
-            height_regularization = torch.true_divide(height_diff, normalization).mean()
-            loss += self.lambda_height_reg * height_regularization
-            loss_dict['height_regularization'] = height_regularization
-
-        loss_dict['loss'] = loss
+        model_out = self.model(coords)
+        model_out_keys = model_out.keys()
+
+        # split back into dataloaders
+        result_mapping = {}
+        idx = 0
+        for k in loader_keys:
+            n_coords = batch[k]['coords'].shape[0]
+            result_mapping[k] = {mk: model_out[mk][idx:idx + n_coords] for mk in model_out_keys}
+            idx += n_coords
+
+        state_dict = {k: {**result_mapping[k], **batch[k]} for k in loader_keys}
+
+        # global state
+        state_dict['all'] = {**model_out, 'coords': coords}
+
+        loss_dict = {}
+        for name, loss_module in self.loss_modules.items():
+            ds_id = loss_module.ds_id
+            if isinstance(ds_id, list):
+                states = [state_dict[i] for i in ds_id]
+                state_keys = states[0].keys()
+                state = {k: torch.cat([s[k] for s in states]) for k in state_keys}
+            else:
+                state = state_dict[ds_id]
+            loss = {name: loss_module(**state)}
+            loss_dict.update(loss)
+        total_loss = sum([self.lambdas[k] * loss_dict[k] for k in loss_dict.keys()])
+
+        loss_dict['loss'] = total_loss
         return loss_dict
 
+    @torch.no_grad()
     def on_train_batch_end(self, outputs, batch, batch_idx) -> None:
-        # update training parameters and log
-        if self.lambda_B > self.lambda_B_end:
-            self.lambda_B *= self.lambda_B_gamma
-        if self.scheduler.get_last_lr()[0] > 5e-5:
-            self.scheduler.step()
-        self.log('Learning Rate', self.scheduler.get_last_lr()[0])
-        self.log('Lambda B', self.lambda_B)
-
+        # update lambda parameters and log
+        for k in self.scheduled_lambdas.keys():
+            param = self.lambdas[k]
+            gamma = self.scheduled_lambdas[k]['gamma']
+            if self.scheduled_lambdas[k]['type'] == 'exponential':
+                if (gamma < 1 and param > self.scheduled_lambdas[k]['end']) or \
+                        (gamma > 1 and param < self.scheduled_lambdas[k]['end']):
+                    new_lambda = param * gamma
+                    param.copy_(new_lambda)
+            if self.scheduled_lambdas[k]['type'] == 'linear':
+                if (gamma > 0 and param < self.scheduled_lambdas[k]['end']) or \
+                        (gamma < 0 and param > self.scheduled_lambdas[k]['end']):
+                    new_lambda = param + gamma
+                    param.copy_(new_lambda)
+            self.log('lambda_' + k, self.lambdas[k])
+        # update learning rate
+        scheduler = self.lr_schedulers()
+        if scheduler.get_last_lr()[0] > self.lr_params['end']:
+            scheduler.step()
+        self.log('Learning Rate', scheduler.get_last_lr()[0])
+
+        assert not torch.isnan(outputs['loss'].mean()), f"Loss is NaN. Check input data and run configuration."
         # log results to WANDB
-        self.log("train/loss", outputs['loss'])
-        self.log("Training Loss", {k: v.mean() for k, v in outputs.items()})
+        self.log("train", {k: v.mean() for k, v in outputs.items()})
 
     @torch.enable_grad()
     def validation_step(self, batch, batch_nb, dataloader_idx):
-        if dataloader_idx == 0:
-            coords = batch
-            coords.requires_grad = True
-
-            # forward step
-            b = self.model(coords)
-
-            jac_matrix = jacobian(b, coords)
-            dBx_dx = jac_matrix[:, 0, 0]
-            dBy_dx = jac_matrix[:, 1, 0]
-            dBz_dx = jac_matrix[:, 2, 0]
-            dBx_dy = jac_matrix[:, 0, 1]
-            dBy_dy = jac_matrix[:, 1, 1]
-            dBz_dy = jac_matrix[:, 2, 1]
-            dBx_dz = jac_matrix[:, 0, 2]
-            dBy_dz = jac_matrix[:, 1, 2]
-            dBz_dz = jac_matrix[:, 2, 2]
-            #
-            rot_x = dBz_dy - dBy_dz
-            rot_y = dBx_dz - dBz_dx
-            rot_z = dBy_dx - dBx_dy
-            #
-            j = torch.stack([rot_x, rot_y, rot_z], -1)
-            div = torch.abs(dBx_dx + dBy_dy + dBz_dz)
-
-            return {'b': b.detach().cpu(), 'j': j.detach().cpu(), 'div': div.detach().cpu()}
-        if dataloader_idx == 1:
-            boundary_coords = batch['coords']
-            b_true = batch['values']
-            b_err = batch['error'] if 'error' in batch else 0
-            boundary_coords.requires_grad = True
-            if self.use_height_mapping:
-                boundary_coords = self.height_mapping_model(boundary_coords, batch['height_ranges'])
-            b = self.model(boundary_coords)
-
-            # compute boundary loss
-            b_diff_error = torch.clip(torch.abs(b - b_true) - b_err, 0)
-            b_diff_error = torch.mean(torch.nansum(b_diff_error.pow(2), -1).pow(0.5))
-            b_diff = torch.abs(b - b_true)
-            b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1).pow(0.5))
-            return {'b_diff_error': b_diff_error.detach().cpu(), 'b_diff': b_diff.detach().cpu()}
-        else:
-            raise NotImplementedError('Validation data loader not supported!')
+        coords = batch['coords']
+        coords.requires_grad = True
+
+        if self.transform_module is not None:
+            if self.validation_mapping[dataloader_idx] in self.transform_module.validation_ds_id:
+                batch = self.transform_module.transform_batch(batch)
+
+        model_out = self.model(coords)
+
+        jac_matrix = model_out['jac_matrix']
+        dBx_dx = jac_matrix[:, 0, 0]
+        dBy_dx = jac_matrix[:, 1, 0]
+        dBz_dx = jac_matrix[:, 2, 0]
+        dBx_dy = jac_matrix[:, 0, 1]
+        dBy_dy = jac_matrix[:, 1, 1]
+        dBz_dy = jac_matrix[:, 2, 1]
+        dBx_dz = jac_matrix[:, 0, 2]
+        dBy_dz = jac_matrix[:, 1, 2]
+        dBz_dz = jac_matrix[:, 2, 2]
+        #
+        rot_x = dBz_dy - dBy_dz
+        rot_y = dBx_dz - dBz_dx
+        rot_z = dBy_dx - dBx_dy
+        #
+        j = torch.stack([rot_x, rot_y, rot_z], -1)
+        div = torch.abs(dBx_dx + dBy_dy + dBz_dz)
+
+        model_out['j'] = j
+        model_out['div'] = div
+
+        return {k: v.detach() for k, v in {**batch, **model_out}.items()}
 
     def validation_epoch_end(self, outputs_list):
-        if len(outputs_list) == 0:
-            return # skip invalid validation steps
-        outputs = outputs_list[0]  # unpack data loader 0
-        # stack magnetic field and unnormalize
-        b = torch.cat([o['b'] for o in outputs]) * self.validation_settings['gauss_per_dB']
-        j = torch.cat([o['j'] for o in outputs]) * self.validation_settings['gauss_per_dB'] / self.validation_settings['Mm_per_ds']
-        div = torch.cat([o['div'] for o in outputs]) * self.validation_settings['gauss_per_dB'] / self.validation_settings['Mm_per_ds']
-
-        norm = b.pow(2).sum(-1).pow(0.5) * j.pow(2).sum(-1).pow(0.5)
-        angle = torch.cross(j, b, dim=-1).pow(2).sum(-1).pow(0.5) / norm
-        sig = torch.asin(torch.clip(angle, -1. + 1e-7, 1. - 1e-7)) * (180 / np.pi)
-        sig = torch.abs(sig)
-        weighted_sig = np.average(sig.cpu().numpy(), weights=j.pow(2).sum(-1).pow(0.5).cpu().numpy())
-        b_norm = b.pow(2).sum(-1).pow(0.5) + 1e-7
-        div_loss = (div / b_norm).mean()
-        ff_loss = torch.cross(j, b, dim=-1).pow(2).sum(-1).pow(0.5) / b_norm
-        ff_loss = ff_loss.mean()
-
-        outputs = outputs_list[1]  # unpack data loader 1
-        b_diff = torch.stack([o['b_diff'] for o in outputs]).mean() * self.validation_settings['gauss_per_dB']
-        b_diff_error = torch.stack([o['b_diff_error'] for o in outputs]).mean() * self.validation_settings['gauss_per_dB']
-
-        b_cube = b.reshape([*self.validation_settings['cube_shape'], 3]).cpu().numpy()
-        self.plot_sample(b_cube)
-
-        self.log("Validation B_diff", b_diff)
-        self.log("Validation B_diff_error", b_diff_error)
-        self.log("Validation DIV", div_loss)
-        self.log("Validation FF", ff_loss)
-        self.log("Validation Sigma", weighted_sig)
-
-        return {'progress_bar': {'b_diff': b_diff, 'div': div_loss, 'ff': ff_loss, 'sigma': weighted_sig},
-                'log': {'val/b_diff': b_diff, 'val/div': div_loss, 'val/ff': ff_loss, 'val/sig': weighted_sig}}
-
-    def plot_sample(self, b, n_samples=10):
-        fig, axs = plt.subplots(3, n_samples, figsize=(n_samples * 4, 12))
-        heights = np.linspace(0, 1, n_samples) ** 2 * (b.shape[2] - 1)  # more samples from lower heights
-        heights = heights.astype(np.int32)
-        for i in range(3):
-            for j, h in enumerate(heights):
-                v_min_max = np.max(np.abs(b[:, :, h, i]))
-                axs[i, j].imshow(b[:, :, h, i].transpose(), cmap='gray', vmin=-v_min_max, vmax=v_min_max,
-                                 origin='lower')
-                axs[i, j].set_axis_off()
-        for j, h in enumerate(heights):
-            axs[0, j].set_title('%.01f' % h)
-        fig.tight_layout()
-        wandb.log({"Slices": fig})
-        plt.close('all')
+        self.validation_outputs = {}  # reset validation outputs
+        if len(outputs_list) == 0 or any([len(o) == 0 for o in outputs_list]):
+            return  # skip invalid validation steps
+
+        for i, outputs in enumerate(outputs_list):
+            out_keys = outputs[0].keys()
+            outputs = {k: torch.cat([o[k] for o in outputs]) for k in out_keys}
+            self.validation_outputs[self.validation_mapping[i]] = outputs
 
     def on_load_checkpoint(self, checkpoint):
-        super().on_load_checkpoint(checkpoint)
-        # update scheduled lambda
-        self.lambda_B *= self.lambda_B_gamma ** checkpoint['global_step']
-        self.lambda_B = max([self.lambda_B, self.lambda_B_end])
-
-
-def calculate_loss(b, coords):
-    jac_matrix = jacobian(b, coords)
-    dBx_dx = jac_matrix[:, 0, 0]
-    dBy_dx = jac_matrix[:, 1, 0]
-    dBz_dx = jac_matrix[:, 2, 0]
-    dBx_dy = jac_matrix[:, 0, 1]
-    dBy_dy = jac_matrix[:, 1, 1]
-    dBz_dy = jac_matrix[:, 2, 1]
-    dBx_dz = jac_matrix[:, 0, 2]
-    dBy_dz = jac_matrix[:, 1, 2]
-    dBz_dz = jac_matrix[:, 2, 2]
-    #
-    rot_x = dBz_dy - dBy_dz
-    rot_y = dBx_dz - dBz_dx
-    rot_z = dBy_dx - dBx_dy
-    #
-    j = torch.stack([rot_x, rot_y, rot_z], -1)
-    jxb = torch.cross(j, b, -1)
-    force_loss = torch.sum(jxb ** 2, dim=-1) / (torch.sum(b ** 2, dim=-1) + 1e-7)
-    divergence_loss = (dBx_dx + dBy_dy + dBz_dz) ** 2
-    return divergence_loss, force_loss
-
-
-def save(save_path, model, data_module, config, height_mapping_model=None):
-    save_state = {'model': model,
-              'cube_shape': data_module.cube_shape,
-              'b_norm': data_module.b_norm,
-              'spatial_norm': data_module.spatial_norm,
-              'meta_data': data_module.meta_data,
-              'config': config,
-              'height_mapping_model': height_mapping_model,
-              'height_mapping': data_module.height_mapping if hasattr(data_module, 'height_mapping') else None,
-              'Mm_per_pixel': data_module.Mm_per_pixel if hasattr(data_module, 'Mm_per_pixel') else None,
-                  }
+        state_dict = checkpoint['state_dict']
+        # keep new lambdas
+        for k, v in self.lambdas.items():
+            if f"lambdas.{k}" not in state_dict:
+                print(f'Add lambda {k}: {v.data}')
+                state_dict[f'lambdas.{k}'] = v
+                continue
+            checkpoint_v = state_dict[f"lambdas.{k}"]
+            if k in self.scheduled_lambdas or checkpoint_v == v:  # skip scheduled lambdas or same values
+                continue
+            print(f'Update lambda {k}: {checkpoint_v} --> {v.data}')
+            state_dict[f'lambdas.{k}'] = v
+        # remove old lambdas
+        remove_keys = []
+        for k, v in state_dict.items():
+            if 'lambdas' in k and k.split('.')[1] not in self.lambdas.keys():
+                print(f'Remove lambda: {k}')
+                remove_keys.append(k)
+        [state_dict.pop(k) for k in remove_keys]
+        # remove old scheduled lambdas
+        remove_keys = []
+        for k, v in state_dict.items():
+            if 'scheduled_lambdas' in k and k.split('.')[1] not in self.scheduled_lambdas.keys():
+                print(f'Remove scheduled lambda: {k}')
+                remove_keys.append(k)
+        [state_dict.pop(k) for k in remove_keys]
+
+        self.load_state_dict(state_dict, strict=False)
+        self.validation_outputs = {}  # reset validation outputs
+
+
+def save(save_path, nf2, data_module, config):
+    save_state = {'model': nf2.model,
+                  'config': config,
+                  'data': data_module.config, }
+    if nf2.transform_module is not None:
+        save_state['transform_module'] = nf2.transform_module
     torch.save(save_state, save_path)
diff --git a/notebooks/NF2_SHARP_extrapolation.ipynb b/notebooks/NF2_SHARP_extrapolation.ipynb
new file mode 100644
index 0000000..c8e4252
--- /dev/null
+++ b/notebooks/NF2_SHARP_extrapolation.ipynb
@@ -0,0 +1,315 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Neural Network Force-Free magnetic field extrapolation - NF2\n",
+        "<img src=\"https://github.com/RobertJaro/NF2/blob/main/images/logo.jpg?raw=true\" width=\"150\" height=\"150\">\n",
+        "\n",
+        "This notebook provides basic NF2 extrapolations of SHARP regions (SDO/HMI). Fill the form below to select your observation and start the extrapolation.\n",
+        "\n",
+        "Run all sections below to start your non-linear force-free extrapolation. Extrapolations take about 6 hours to complete in Google Colab. All results are logged to Weights And Biases (https://wandb.ai). If you want to continue an interrupted run, fill the WandB section.\n",
+        "\n",
+        "GitHub Page: https://github.com/RobertJaro/NF2\n"
+      ],
+      "metadata": {
+        "id": "v_e2_34GKe4s"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Install NF2"
+      ],
+      "metadata": {
+        "id": "LjKmZKM2dKZz"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qBle4tJeJP6O"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/RobertJaro/NF2@v0.3.0\n",
+        "\n",
+        "from nf2.data.download import download_SHARP_series\n",
+        "from nf2.extrapolate import run\n",
+        "# util functions\n",
+        "from datetime import datetime"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Data Download"
+      ],
+      "metadata": {
+        "id": "VYJL8JtUNXwa"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Downloading data requires an active registration at JSOC. http://jsoc.stanford.edu/ajax/register_email.html"
+      ],
+      "metadata": {
+        "id": "DNa_LMbPNalC"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wJRSB1NV--Bb"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Download Credentials\n",
+        "jsoc_email = 'robert.jarolim@uni-graz.at' #@param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2zVLpXBf2oQN",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Date of observation\n",
+        "year = 2017 #@param {type:\"integer\"}\n",
+        "month = 9 #@param {type:\"integer\"}\n",
+        "day = 6 #@param {type:\"integer\"}\n",
+        "hour = 8 #@param {type:\"integer\"}\n",
+        "minute = 36 #@param {type:\"number\"}\n",
+        "\n",
+        "date = datetime(year, month, day, hour, minute)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "nDUBoZ3-86en",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Active Region\n",
+        "#@markdown specify either the NOAA or SHARP number of the active region\n",
+        "noaa_num = 12673 #@param {type:\"number\"}\n",
+        "sharp_num = None #@param {type:\"number\"}\n",
+        "\n",
+        "download_dir = 'AR_12673' #@param {type:\"string\"}\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "download_result = download_SHARP_series(download_dir=download_dir,\n",
+        "                                        email=jsoc_email, t_start=date,\n",
+        "                                        noaa_num=noaa_num, sharp_num=sharp_num)\n",
+        "download_result"
+      ],
+      "metadata": {
+        "id": "kPaLqvylOR7_"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Extrapolation"
+      ],
+      "metadata": {
+        "id": "VgzWQC98RECn"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "iJaeLqddRY8e"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Storage path\n",
+        "base_path = 'extrapolation_AR12673' #@param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "QCMzy2_vR4YM"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Data configuration\n",
+        "#@markdown select directory for saving extrapolation results\n",
+        "base_path = 'extrapolation_AR12673' #@param {type:\"string\"}\n",
+        "\n",
+        "#@markdown insert downloaded SHARP files\n",
+        "br_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Br.fits' #@param {type:\"string\"}\n",
+        "bt_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Bt.fits' #@param {type:\"string\"}\n",
+        "bp_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Bp.fits' #@param {type:\"string\"}\n",
+        "\n",
+        "#@markdown insert downloaded SHARP error files (optional)\n",
+        "use_error_files = True #@param {type:\"boolean\"}\n",
+        "br_error_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Br_err.fits' #@param {type:\"string\"}\n",
+        "bt_error_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Bt_err.fits' #@param {type:\"string\"}\n",
+        "bp_error_path = '/content/AR_12673/hmi.sharp_cea_720s.7115.20170906_083600_TAI.Bp_err.fits' #@param {type:\"string\"}\n",
+        "\n",
+        "data_config = {\n",
+        "  'type': 'fits',\n",
+        "  'slices': [{\n",
+        "      'fits_path' : {\n",
+        "        'Br': br_path,\n",
+        "        'Bt': bt_path,\n",
+        "        'Bp': bp_path\n",
+        "      }\n",
+        "  }],\n",
+        "  'iterations': 10000\n",
+        "}\n",
+        "\n",
+        "if use_error_files:\n",
+        "  data_config['slices'][0]['error_path'] = {'Br_err': br_error_path,\n",
+        "                                         'Bt_err': bt_error_path,\n",
+        "                                         'Bp_err': bp_error_path}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Continue training from WandB\n",
+        "#@markdown Enter your credentials if you want to continue an inturrpted run. This will restore the latest model and continue the training.\n",
+        "continue_training = False #@param {type:\"boolean\"}\n",
+        "id = \"\" #@param {type:\"string\"}\n",
+        "entity = \"\" #@param {type:\"string\"}\n",
+        "project = \"\" #@param {type:\"string\"}\n",
+        "\n",
+        "if continue_training:\n",
+        "  logging_config = {'id': id, 'entity': entity, 'project': project}\n",
+        "else:\n",
+        "  logging_config = {}"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "Avh_MFrzdLVr"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Advanced model settings\n",
+        "use_custom_settings = False #@param {type:\"boolean\"}\n",
+        "#@markdown Set `vector_potential` if you want to run a divergence-free extrapolation\n",
+        "vector_potential = False #@param {type:\"boolean\"}\n",
+        "#@markdown Change `dim` if you want to run a divergence-free extrapolation\n",
+        "dim = 256 #@param {type:\"number\"}\n",
+        "\n",
+        "if use_custom_settings:\n",
+        "  model_config = {'type': 'vector_potential' if vector_potential else 'b',\n",
+        "                  'dim': dim}\n",
+        "else:\n",
+        "  model_config = None"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "DIF8N5VcWOg0"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Advanced training settings\n",
+        "use_custom_settings = False #@param {type:\"boolean\"}\n",
+        "#@markdown Change the weighting of the force-free condition\n",
+        "lambda_ff = 0.1 #@param {type:\"number\"}\n",
+        "#@markdown Change the weighting of the divergence-free condition\n",
+        "lambda_div = 0.1 #@param {type:\"number\"}\n",
+        "#@markdown Change the weighting of the boundary condition\n",
+        "lambda_boundary = 1.0 #@param {type:\"number\"}\n",
+        "\n",
+        "\n",
+        "\n",
+        "if use_custom_settings:\n",
+        "  boundary_loss = {'type': 'boundary', 'name': 'boundary', 'lambda': lambda_boundary,\n",
+        "                 'ds_id': ['boundary_01', 'potential']}\n",
+        "  force_free_loss = {'type': 'force_free', 'lambda': lambda_ff}\n",
+        "  divergence_free_loss = {'type': 'divergence', 'lambda': lambda_div}\n",
+        "\n",
+        "  training_config = {'loss_config': [boundary_loss, force_free_loss, divergence_free_loss]}\n",
+        "else:\n",
+        "  training_config = {}"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "fQRYPHNfXR_g"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Start Extrapolation\n",
+        "For the first run you will be redirected to the WandB login. Login with your existing user or register a new user to log your results online."
+      ],
+      "metadata": {
+        "id": "7gwv-nIXbpW8"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "run(base_path, data=data_config,\n",
+        "    logging=logging_config,\n",
+        "    training=training_config,\n",
+        "    model=model_config)"
+      ],
+      "metadata": {
+        "id": "MMROBK3DRt2A"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "7wGGAZZWewz_"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file
diff --git a/notebooks/NF2_convert.ipynb b/notebooks/NF2_convert.ipynb
new file mode 100644
index 0000000..a59dedf
--- /dev/null
+++ b/notebooks/NF2_convert.ipynb
@@ -0,0 +1,206 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Neural Network Force-Free magnetic field extrapolation - NF2\n",
+        "<img src=\"https://github.com/RobertJaro/NF2/blob/main/images/logo.jpg?raw=true\" width=\"150\" height=\"150\">\n",
+        "\n",
+        "This notebook provides conversions of NF2 files to common output formats.\n",
+        "\n",
+        "GitHub Page: https://github.com/RobertJaro/NF2\n"
+      ],
+      "metadata": {
+        "id": "v_e2_34GKe4s"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Install NF2"
+      ],
+      "metadata": {
+        "id": "LjKmZKM2dKZz"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qBle4tJeJP6O"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/RobertJaro/NF2@v0.3.0\n",
+        "\n",
+        "from nf2.evaluation.output import CartesianOutput"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Load NF2 result"
+      ],
+      "metadata": {
+        "id": "qRxa8Wq7LBCx"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Upload your NF2 file and specify path\n",
+        "nf2_path = 'extrapolation_result.nf2' #@param {type:\"string\"}"
+      ],
+      "metadata": {
+        "id": "ZBlqkLx3JTy_"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Custom resolution and height (Mm)\n",
+        "Mm_per_pixel = 0.72 #@param {type:\"number\"}\n",
+        "max_height = 80 #@param {type:\"number\"}"
+      ],
+      "metadata": {
+        "id": "OAXFjg22KM_a"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The NF2 `CartesianOutput` loads all information from the NF2 file. It can be used to load custom coordinates or the full data cube as secified in the NF2 file. If no `Mm_per_pixel` and `height_range` are specified, the default values from the extrapolation are used. For full resolution SHARPs this can require a lot of memory and computing time."
+      ],
+      "metadata": {
+        "id": "RhzkCQILQ_7I"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "nf2_out = CartesianOutput(nf2_path)\n",
+        "output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=[0, max_height], progress=True)"
+      ],
+      "metadata": {
+        "id": "7wGGAZZWewz_"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "print('NF2 Outputs:')\n",
+        "for k, v in output.items():\n",
+        "  print(f'{k} : {v.shape if hasattr(v, \"shape\") else v}')"
+      ],
+      "metadata": {
+        "id": "sqk7FqtuKrRe"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Convert NF2 to VTK"
+      ],
+      "metadata": {
+        "id": "ReP6XcloMxxj"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Convert NF2 file to VTK for visualization in Paraview. This requires the installation of TVTK (mayavi)."
+      ],
+      "metadata": {
+        "id": "ifaXxJQmNVXC"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install mayavi"
+      ],
+      "metadata": {
+        "id": "UOCOzk8xNBqs"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!nf2-to-vtk --nf2_path \"extrapolation_result.nf2\" --out_path \"extrapolation_result.vtk\" --Mm_per_pixel 0.72"
+      ],
+      "metadata": {
+        "id": "liD0hiTiMYuU"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Convert NF2 to other formats"
+      ],
+      "metadata": {
+        "id": "2FSnk4u0ODLG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!nf2-to-npy --nf2_path \"extrapolation_result.nf2\" --out_path \"extrapolation_result.npy\" --Mm_per_pixel 0.72"
+      ],
+      "metadata": {
+        "id": "Q-KgGiGOOGxI"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!nf2-to-hdf5 --nf2_path \"extrapolation_result.nf2\" --out_path \"extrapolation_result.hdf5\" --Mm_per_pixel 0.72"
+      ],
+      "metadata": {
+        "id": "mh2yFO-cO4yL"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!nf2-to-fits --nf2_path \"extrapolation_result.nf2\" --out_path \"extrapolation_result.fits\" --Mm_per_pixel 0.72"
+      ],
+      "metadata": {
+        "id": "r8q8xULwPCQ7"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 7b7c1dc..34ca5b3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -11,4 +11,5 @@ setuptools>=49.6.0
 wandb>=0.13
 lightning==1.9.3
 pytorch_lightning==1.9.3
-mayavi
\ No newline at end of file
+mayavi
+pfsspy
\ No newline at end of file
diff --git a/scripts/download.sh b/scripts/download.sh
index 022d2c6..7a83600 100644
--- a/scripts/download.sh
+++ b/scripts/download.sh
@@ -1,3 +1,5 @@
 module load python/pytorch-1.6.0
 cd /beegfs/home/robert.jarolim/projects/pub_NF2
-python3 -m nf2.data.download_range --download_dir /gpfs/gpfs0/robert.jarolim/data/nf2/7115 --email robert.jarolim@uni-graz.at --harpnum 7115 --t_start 2017-09-02T00:00:00 --duration 6d
+#python3 -m nf2.data.download_sharp --download_dir /gpfs/gpfs0/robert.jarolim/data/nf2/7310 --email robert.jarolim@uni-graz.at --harpnum 7310 --t_start 2018-09-30T09:24:00 --series "sharp_cea_720s"
+#python3 -m nf2.data.download_range --download_dir /gpfs/gpfs0/robert.jarolim/data/nf2/7115 --email robert.jarolim@uni-graz.at --harpnum 7115 --t_start 2017-09-02T00:00:00 --duration 6d
+python3 -m nf2.data.download_sharp --download_dir /gpfs/gpfs0/robert.jarolim/data/nf2/az_377 --email robert.jarolim@uni-graz.at --harpnum 377 --t_start 2011-02-15T00:00:00 --series "sharp_720s" --segments 'field, inclination, azimuth, disambig'
\ No newline at end of file
diff --git a/scripts/evaluate_series.sh b/scripts/evaluate_series.sh
index 6a1cd43..b1e08a1 100644
--- a/scripts/evaluate_series.sh
+++ b/scripts/evaluate_series.sh
@@ -1,3 +1,16 @@
+#!/bin/bash
+#SBATCH --job-name=nf2                       # Job name
+#SBATCH --partition=gpu                      # Queue name
+#SBATCH --nodes=1                            # Run all processes on a single node
+#SBATCH --ntasks-per-node=8                          # Run 12 tasks
+#SBATCH --mem=64000                          # Job memory request in Megabytes
+#SBATCH --gpus=2                             # Number of GPUs
+#SBATCH --time=24:00:00                      # Time limit hrs:min:sec or dd-hrs:min:sec
+#SBATCH --output=/gpfs/gpfs0/robert.jarolim/nf2/logs/nf2_train_%j.log     # Standard output and error log
+
+
 module load python/pytorch-1.6.0
 cd /beegfs/home/robert.jarolim/projects/pub_NF2
-python3 -i -m nf2.evaluation.series /gpfs/gpfs0/robert.jarolim/nf2/7115/series --result_path /gpfs/gpfs0/robert.jarolim/nf2/7115/evaluation --strides 1 --flare_list /gpfs/gpfs0/robert.jarolim/data/goes_flares_integrated.csv
+#python3 -i -m nf2.evaluation.series /gpfs/gpfs0/robert.jarolim/nf2/401/series --result_path /gpfs/gpfs0/robert.jarolim/nf2/401/evaluation --strides 2 --add_flares
+#python3 -i -m nf2.evaluation.series /gpfs/gpfs0/robert.jarolim/nf2/6975/series --result_path /gpfs/gpfs0/robert.jarolim/nf2/6975/evaluation --strides 4 --add_flares
+python3 -i -m nf2.evaluation.series /gpfs/gpfs0/robert.jarolim/nf2/8088/series --result_path /gpfs/gpfs0/robert.jarolim/nf2/8088/evaluation2 --strides 1 --add_flares
diff --git a/setup.py b/setup.py
index 1d6872b..1c3e8be 100644
--- a/setup.py
+++ b/setup.py
@@ -11,15 +11,18 @@
     author_email='',
     description='Neural Network Force Free magnetic field extrapolation',
     install_requires=['torch>=1.8', 'sunpy[all]>=3.0', 'scikit-image', 'scikit-learn', 'tqdm',
-                      'numpy', 'matplotlib', 'astropy', 'drms', 'wandb>=0.13', 'lightning==1.9.3', 'pytorch_lightning==1.9.3'],
+                      'numpy', 'matplotlib', 'astropy', 'drms', 'wandb>=0.13', 'lightning==1.9.3', 'pytorch_lightning==1.9.3',
+                      'pfsspy'],
     entry_points={
         'console_scripts': [
             'nf2-extrapolate = nf2.extrapolate:main',
             'nf2-extrapolate-series = nf2.extrapolate_series:main',
             'nf2-noaa-to-sharp = nf2.data.noaa_to_sharp:main',
             'nf2-download = nf2.data.download_range:main',
-            'nf2-to-vtk = nf2.evaluation.nf2_to_vtk:main',
-            'nf2-to-npy = nf2.evaluation.nf2_to_vtk:main',
-            'nf2-to-fits = nf2.evaluation.nf2_to_fits:main',
+            'nf2-to-vtk = nf2.convert.nf2_to_vtk:main',
+            'nf2-to-npy = nf2.convert.nf2_to_npy:main',
+            'nf2-to-fits = nf2.convert.nf2_to_fits:main',
+            'nf2-to-hdf5 = nf2.convert.nf2_to_hdf5:main',
+            'nf2-convert-series = nf2.convert.convert_series:main',
         ]}
 )


Test patch:
diff --git a/config/fd_2173_potential_test.json b/config/fd_2173_potential_test.json
new file mode 100644
index 0000000..c79e9ad
--- /dev/null
+++ b/config/fd_2173_potential_test.json
@@ -0,0 +1,60 @@
+{
+  "base_path": "/glade/work/rjarolim/global/2173_potential_test_v1",
+  "logging": {
+    "wandb_entity": "robert_jarolim",
+    "wandb_project": "global_HAO",
+    "wandb_name": "2173_potential_test",
+    "wandb_id": null
+  },
+  "data": {
+    "type": "potential_test",
+    "synoptic_files": {
+      "Br": "/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits"
+    },
+    "boundary": {
+      "type": "open"
+    },
+    "height": 1.3,
+    "b_norm": 2000,
+    "batch_size": {
+      "boundary": 8192,
+      "random": 16384
+    },
+    "work_directory": "/glade/derecho/scratch/rjarolim/nf2/2173_potential_test_v1",
+    "num_workers": 8
+  },
+  "model": {
+    "dim": 512,
+    "use_vector_potential": true
+  },
+  "training": {
+    "epochs": 1e4,
+    "lambda_b": {
+      "start": 1e2,
+      "end": 10,
+      "iterations": 2e5
+    },
+    "lambda_force_free": 1e-2,
+    "lambda_radial": 0,
+    "lambda_energy_gradient": 1e-2,
+    "check_val_every_n_epoch": 100
+  },
+  "plot": [
+    {
+      "type": "slices",
+      "name": "subframe",
+      "longitude_range": [2.356194490192345, 3.141592653589793],
+      "latitude_range": [1.2566370614359172, 1.8849555921538759],
+      "radius_range": [1.0, 1.2],
+      "n_slices": 10,
+      "plot_currents": true
+    },
+    {
+      "type": "slices",
+      "name": "full-disk",
+      "radius_range": [1.0, 1.3],
+      "n_slices": 5,
+      "plot_currents": true
+    }
+  ]
+}
\ No newline at end of file