Repository: RobertJaro/NF2 (language: python)

Issue:
pull v0.3.0 changes



Discussion / hints:

Fix patch:
diff --git a/config/2106_AR.json b/config/2106_AR.json
deleted file mode 100644
index 7de5c8e..0000000
--- a/config/2106_AR.json
+++ /dev/null
@@ -1,78 +0,0 @@
-{
-  "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
deleted file mode 100644
index 631155f..0000000
--- a/config/2173.json
+++ /dev/null
@@ -1,86 +0,0 @@
-{
-  "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
deleted file mode 100644
index 4bc2878..0000000
--- a/config/2173_full_disc.json
+++ /dev/null
@@ -1,81 +0,0 @@
-{
-  "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
deleted file mode 100644
index 9a8ab56..0000000
--- a/config/2173_submap.json
+++ /dev/null
@@ -1,101 +0,0 @@
-{
-  "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
deleted file mode 100644
index bceefb6..0000000
--- a/config/2173_synoptic_ref.json
+++ /dev/null
@@ -1,75 +0,0 @@
-{
-  "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
deleted file mode 100644
index a860981..0000000
--- a/config/377_az.json
+++ /dev/null
@@ -1,58 +0,0 @@
-{
-  "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_16.json b/config/analytical_0_16.json
deleted file mode 100644
index f585517..0000000
--- a/config/analytical_0_16.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_0_16",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "0_16",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0, 16]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_0_16",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_0_32.json b/config/analytical_0_32.json
deleted file mode 100644
index 5cdfd17..0000000
--- a/config/analytical_0_32.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_0_32",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "0_32",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0, 32]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_0_32",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_0_8.json b/config/analytical_0_8.json
deleted file mode 100644
index 6a9b496..0000000
--- a/config/analytical_0_8.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_0_8",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "0_8",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0, 8]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_0_8",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_0_8_16.json b/config/analytical_0_8_16.json
deleted file mode 100644
index ad830f7..0000000
--- a/config/analytical_0_8_16.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_0_8_16",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "0_8_16",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0, 8, 16]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_0_8_16",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_0_8_16_32.json b/config/analytical_0_8_16_32.json
deleted file mode 100644
index b5ddb90..0000000
--- a/config/analytical_0_8_16_32.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/glade/work/rjarolim/nf2/analytical/analytical_0_8_16_32",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "0_8_16_32",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0, 8, 16, 32]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/glade/work/rjarolim/nf2/analytical/analytical_0_8_16_32",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_1tau.json b/config/analytical_1tau.json
deleted file mode 100644
index 70541e0..0000000
--- a/config/analytical_1tau.json
+++ /dev/null
@@ -1,43 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_l1e-1",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "muram_l1e-1",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, 1, 2, 3, 4],
-    "height_mapping": {
-      "z":  [0.000, 1.014, 1.804, 2.663, 4.851],
-      "z_min": [0.000, 0.000, 0.000, 0.000, 0.000],
-      "z_max": [0.000, 1.658,12.013,19.495,67.942]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "open"},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_l1e-1",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "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}
-  }
-}
\ No newline at end of file
diff --git a/config/analytical_2tau.json b/config/analytical_2tau.json
deleted file mode 100644
index 6d2b990..0000000
--- a/config/analytical_2tau.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_2tau",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "2tau",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": false, "tau_surfaces" : [0, 8]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_2tau",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_2tau_Bz.json b/config/analytical_2tau_Bz.json
deleted file mode 100644
index 20e2ef4..0000000
--- a/config/analytical_2tau_Bz.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_2tau_Bz",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "2tau_Bz",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": true, "tau_surfaces" : [0, 8]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_2tau_Bz",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_3tau.json b/config/analytical_3tau.json
deleted file mode 100644
index 0425b8d..0000000
--- a/config/analytical_3tau.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_3tau",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "3tau",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": false, "tau_surfaces" : [0, 8, 16]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_3tau",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_3tau_Bz.json b/config/analytical_3tau_Bz.json
deleted file mode 100644
index 73bdbee..0000000
--- a/config/analytical_3tau_Bz.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_3tau_Bz",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "3tau_Bz",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": true, "tau_surfaces" : [0, 8, 16]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_3tau_Bz",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_4tau.json b/config/analytical_4tau.json
deleted file mode 100644
index b2a12a4..0000000
--- a/config/analytical_4tau.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_4tau",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "4tau",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": false, "tau_surfaces" : [0, 8, 16, 32]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_4tau",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_4tau_Bz.json b/config/analytical_4tau_Bz.json
deleted file mode 100644
index b4d43f8..0000000
--- a/config/analytical_4tau_Bz.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_realistic_bz_v11",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "4tau_Bz",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "tau", "use_LOS": true, "tau_surfaces" : [0, 8, 16, 32]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_4tau_Bz",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_full.json b/config/analytical_full.json
deleted file mode 100644
index 10da0ea..0000000
--- a/config/analytical_full.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_full",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "full",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "full"},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_full",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_open.json b/config/analytical_open.json
deleted file mode 100644
index d986ab5..0000000
--- a/config/analytical_open.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_open_v2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "open",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "open", "height_slices": [0]},
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_open",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/analytical_potential.json b/config/analytical_potential.json
deleted file mode 100644
index b2ac3b5..0000000
--- a/config/analytical_potential.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/analytical_potential",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "potential",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "analytical",
-    "case": 1,
-    "boundary": {"type":  "potential", "strides": 1 },
-    "height": 64,
-    "b_norm": 300,
-    "spatial_norm": 64,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/analytical_potential",
-    "num_workers": 8,
-    "Mm_per_pixel": 1
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "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/demo_submap.json b/config/demo_submap.json
deleted file mode 100644
index 6274b44..0000000
--- a/config/demo_submap.json
+++ /dev/null
@@ -1,103 +0,0 @@
-{
-  "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
deleted file mode 100644
index 3c72ec4..0000000
--- a/config/demo_synoptic.json
+++ /dev/null
@@ -1,66 +0,0 @@
-{
-  "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
deleted file mode 100644
index 78f1794..0000000
--- a/config/fd_2106.json
+++ /dev/null
@@ -1,69 +0,0 @@
-{
-  "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
deleted file mode 100644
index 98e60c2..0000000
--- a/config/fd_2106_subframe_vp.json
+++ /dev/null
@@ -1,71 +0,0 @@
-{
-  "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
deleted file mode 100644
index 06468fc..0000000
--- a/config/fd_2106_vp.json
+++ /dev/null
@@ -1,73 +0,0 @@
-{
-  "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
deleted file mode 100644
index 9b2b8c6..0000000
--- a/config/fd_2154.json
+++ /dev/null
@@ -1,63 +0,0 @@
-{
-  "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
deleted file mode 100644
index 9664e86..0000000
--- a/config/fd_2154_az.json
+++ /dev/null
@@ -1,62 +0,0 @@
-{
-  "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
deleted file mode 100644
index c6e044a..0000000
--- a/config/fd_2154_vp.json
+++ /dev/null
@@ -1,61 +0,0 @@
-{
-  "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
deleted file mode 100644
index 748da21..0000000
--- a/config/fd_2154_vp_series.json
+++ /dev/null
@@ -1,52 +0,0 @@
-{
-  "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
deleted file mode 100644
index 88fd765..0000000
--- a/config/fd_2173.json
+++ /dev/null
@@ -1,73 +0,0 @@
-{
-  "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
deleted file mode 100644
index 531e698..0000000
--- a/config/fd_2173_synoptic.json
+++ /dev/null
@@ -1,62 +0,0 @@
-{
-  "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
deleted file mode 100644
index 0243ab3..0000000
--- a/config/fd_2173_vp.json
+++ /dev/null
@@ -1,67 +0,0 @@
-{
-  "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
deleted file mode 100644
index 660d310..0000000
--- a/config/fd_2173_vp_series.json
+++ /dev/null
@@ -1,57 +0,0 @@
-{
-  "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
deleted file mode 100644
index fcecefb..0000000
--- a/config/fd_2213_synoptic.json
+++ /dev/null
@@ -1,60 +0,0 @@
-{
-  "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_377.json b/config/hmi_377.json
deleted file mode 100644
index 1b208ad..0000000
--- a/config/hmi_377.json
+++ /dev/null
@@ -1,47 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_377",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_377",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "data_path": [
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br_err.fits"
-    ],
-    "slice": [66,658,9,377],
-    "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": "/tmp/nf2/hmi_377",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "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}
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_377_sampled.json b/config/hmi_377_sampled.json
deleted file mode 100644
index a24c0ce..0000000
--- a/config/hmi_377_sampled.json
+++ /dev/null
@@ -1,47 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_377_sampled",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_377_sampled",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "data_path": [
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bp_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Bt_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/series/hmi.sharp_cea_720s.377.20110212_000000_TAI.Br_err.fits"
-    ],
-    "slice": [66,658,9,377],
-    "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":  1024, "random":  256},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/hmi_377_sampled",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1, "end": 1, "iterations" : 1},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 1e-3, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_377_series.json b/config/hmi_377_series.json
deleted file mode 100644
index 8245e26..0000000
--- a/config/hmi_377_series.json
+++ /dev/null
@@ -1,40 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_377_v2",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_377/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_series_377",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/series",
-    "slice": [66,658,9,377],
-    "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": "/tmp/nf2/hmi_series_377",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_7115.json b/config/hmi_7115.json
deleted file mode 100644
index 91004ee..0000000
--- a/config/hmi_7115.json
+++ /dev/null
@@ -1,44 +0,0 @@
-{
-  "base_path": "/Users/robert/PycharmProjects/NF2/test",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_7115",
-    "wandb_id": null
-  },
-
-  "data": {
-    "type": "sharp",
-    "data_path": [
-    "/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,
-    "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": "/tmp/nf2/hmi_7115",
-    "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}
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_7115_series.json b/config/hmi_7115_series.json
deleted file mode 100644
index baf911b..0000000
--- a/config/hmi_7115_series.json
+++ /dev/null
@@ -1,37 +0,0 @@
-{
-  "base_path": "<<YOUR PATH>>",
-  "meta_path": "<<Initial RUN>>/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "<<YOUR Entity>>",
-    "wandb_project": "nf2",
-    "wandb_name": "hmi_series_7115",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "<<DIR>>",
-    "bin": 2,
-    "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": "/tmp/nf2/hmi_series_7115",
-    "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
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_7115_series_tiny.json b/config/hmi_7115_series_tiny.json
deleted file mode 100644
index 61f3b37..0000000
--- a/config/hmi_7115_series_tiny.json
+++ /dev/null
@@ -1,39 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_7115",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_series_7115",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/7115",
-    "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": "/tmp/nf2/hmi_series_7115",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_7115_tiny.json b/config/hmi_7115_tiny.json
deleted file mode 100644
index 4d18121..0000000
--- a/config/hmi_7115_tiny.json
+++ /dev/null
@@ -1,46 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115_tiny",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_7115_tiny",
-    "wandb_id": null
-  },
-
-  "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"
-    ],
-    "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": "/tmp/nf2/hmi_7115_tiny",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 64,
-    "use_height_mapping": false
-  },
-
-  "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}
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_7115_tiny_series.json b/config/hmi_7115_tiny_series.json
deleted file mode 100644
index 45a6a6c..0000000
--- a/config/hmi_7115_tiny_series.json
+++ /dev/null
@@ -1,39 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115_tiny_series",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115_tiny/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_7115_tiny_series",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/7115",
-    "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": "/tmp/nf2/hmi_7115_tiny_series",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 64,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_8684.json b/config/hmi_8684.json
deleted file mode 100644
index 1bd2a0e..0000000
--- a/config/hmi_8684.json
+++ /dev/null
@@ -1,46 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_8684",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "sst_campaign",
-    "wandb_name": "hmi_8684",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "data_path": [
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_TAI.Bp.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_TAI.Bp_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_TAI.Bt.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_TAI.Bt_err.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_TAI.Br.fits",
-    "/gpfs/gpfs0/robert.jarolim/data/nf2/8684/hmi.sharp_cea_720s_nrt.8684.20230731_060000_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": "/tmp/nf2/hmi_8684",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "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}
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_series_377.json b/config/hmi_series_377.json
deleted file mode 100644
index 8245e26..0000000
--- a/config/hmi_series_377.json
+++ /dev/null
@@ -1,40 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_377_v2",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_377/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_series_377",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/series",
-    "slice": [66,658,9,377],
-    "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": "/tmp/nf2/hmi_series_377",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_series_7115.json b/config/hmi_series_7115.json
deleted file mode 100644
index 9b369ac..0000000
--- a/config/hmi_series_7115.json
+++ /dev/null
@@ -1,39 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_7115",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_series_7115",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/hmi_7115",
-    "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": "/tmp/nf2/hmi_series_7115",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/hmi_series_7115_tiny.json b/config/hmi_series_7115_tiny.json
deleted file mode 100644
index e91cc0f..0000000
--- a/config/hmi_series_7115_tiny.json
+++ /dev/null
@@ -1,39 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_series_7115",
-  "meta_path": "/gpfs/gpfs0/robert.jarolim/multi_height/hmi_7115/extrapolation_result.nf2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "multi-height",
-    "wandb_name": "hmi_series_7115",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "sharp",
-    "paths": "/gpfs/gpfs0/robert.jarolim/data/nf2/series",
-    "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": "/tmp/nf2/hmi_series_7115",
-    "num_workers": 12
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": 1,
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lr_params": 5e-4,
-    "check_val_every_n_epoch": 100
-  }
-}
\ No newline at end of file
diff --git a/config/muram_2tau.json b/config/muram_2tau.json
deleted file mode 100644
index ecb4f97..0000000
--- a/config/muram_2tau.json
+++ /dev/null
@@ -1,43 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_2tau",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, -2],
-    "height_mapping": {
-      "z":  [0.0, 3.0],
-      "z_min": [0.0, 0.0],
-      "z_max": [0.0, 30.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_2tau",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "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}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_2tau_Bxyz_split.json b/config/muram_2tau_Bxyz_split.json
deleted file mode 100644
index a405efd..0000000
--- a/config/muram_2tau_Bxyz_split.json
+++ /dev/null
@@ -1,46 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bxyz_split_v2",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_2tau_Bxyz_split",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, -3, -2],
-    "components": [[0, 1, 2], [0, 1], [2]],
-    "height_mapping": {
-      "z":  [0.0, 2.0, 3.0],
-      "z_min": [0.0, 0.0, 0.0],
-      "z_max": [0.0, 20.0, 30.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_2tau_Bxyz_split",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true,
-    "use_vector_potential": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_min_energy_nans": 0,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_2tau_Bz.json b/config/muram_2tau_Bz.json
deleted file mode 100644
index 682f119..0000000
--- a/config/muram_2tau_Bz.json
+++ /dev/null
@@ -1,45 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_Bz_2epochs",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_2tau_Bz",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, -2],
-    "height_mapping": {
-      "z":  [0.0, 3.0],
-      "z_min": [0.0, 0.0],
-      "z_max": [0.0, 30.0]
-    },
-    "use_bz": true,
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_2tau_Bz_2epochs",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_min_energy_nans": 1e-2,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_2tau_vp.json b/config/muram_2tau_vp.json
deleted file mode 100644
index 202929d..0000000
--- a/config/muram_2tau_vp.json
+++ /dev/null
@@ -1,44 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_2tau_vp",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_2tau_vp",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, -2],
-    "height_mapping": {
-      "z":  [0.0, 3.0],
-      "z_min": [0.0, 0.0],
-      "z_max": [0.0, 30.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_2tau_vp",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true,
-    "use_vector_potential": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 0,
-    "lambda_ff": 1e-1,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_extrapolation.json b/config/muram_extrapolation.json
deleted file mode 100644
index 05f2c47..0000000
--- a/config/muram_extrapolation.json
+++ /dev/null
@@ -1,40 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_extrapolation",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0],
-    "height_mapping": {
-      "z":  [0.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "open"},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_extrapolation",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_extrapolation_pf.json b/config/muram_extrapolation_pf.json
deleted file mode 100644
index 4a1d6db..0000000
--- a/config/muram_extrapolation_pf.json
+++ /dev/null
@@ -1,41 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation_pf",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_extrapolation_pf",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0],
-    "height_mapping": {
-      "z":  [0.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides": 4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_extrapolation_pf",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false,
-    "use_vector_potential": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_extrapolation_vp.json b/config/muram_extrapolation_vp.json
deleted file mode 100644
index c9c9e82..0000000
--- a/config/muram_extrapolation_vp.json
+++ /dev/null
@@ -1,41 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_extrapolation_vp",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_extrapolation_vp",
-    "wandb_id": null
-  },
-  
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0],
-    "height_mapping": {
-      "z":  [0.0]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides": 4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_extrapolation_pf",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false,
-    "use_vector_potential": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_fixed.json b/config/muram_fixed.json
deleted file mode 100644
index 6cb9b28..0000000
--- a/config/muram_fixed.json
+++ /dev/null
@@ -1,41 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_fixed",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_fixed",
-    "wandb_id": null
-  },
-
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, 1, 2, 3, 4, 5],
-    "height_mapping": {
-      "z":   [0.000, 1.018, 1.381, 1.806, 2.685, 4.854]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_fixed",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "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}
-  }
-}
\ No newline at end of file
diff --git a/config/muram_ideal.json b/config/muram_ideal.json
deleted file mode 100644
index 9fbadfc..0000000
--- a/config/muram_ideal.json
+++ /dev/null
@@ -1,43 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/muram_ideal",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "muram_comparison",
-    "wandb_name": "muram_ideal",
-    "wandb_id": null
-  },
-
-  "data": {
-    "type": "numpy",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/muram_slices.npy",
-    "slices": [0, 1, 2, 3, 4, 5],
-    "height_mapping": {
-      "z":   [0.630, 1.018, 1.381, 1.806, 2.685, 4.854],
-      "z_min": [0.000, 0.000, 0.000, 0.000, 0.000, 0.000],
-      "z_max": [0.960, 1.579, 3.882,  11.622, 19.378, 59.592]
-    },
-    "Mm_per_pixel": 0.384,
-    "boundary": {"type":  "potential", "strides": 4},
-    "height": 150,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/muram_ideal",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "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}
-  }
-}
\ No newline at end of file
diff --git a/config/solis_401_chromospheric.json b/config/solis_401_chromospheric.json
deleted file mode 100644
index 6a9a5f1..0000000
--- a/config/solis_401_chromospheric.json
+++ /dev/null
@@ -1,42 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/solis_401_chromospheric_pf_v3",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "solis",
-    "wandb_name": "solis_401_chromospheric_pf"
-  },
-  
-  "data": {
-    "type": "solis",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/401.npy",
-    "height_mapping": {
-      "z":  [0.000, 5.55],
-      "z_min": [0.000, 0.000],
-      "z_max": [0.000, 27.778]
-    },
-    "Mm_per_pixel": 0.36,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 320,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/solis_401_chromospheric_pf_v3",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 1e5},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_min_energy_nans": 1e-2,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/solis_401_photospheric.json b/config/solis_401_photospheric.json
deleted file mode 100644
index f7e653d..0000000
--- a/config/solis_401_photospheric.json
+++ /dev/null
@@ -1,40 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/solis_401_photospheric_v3",
-  "logging": {
-    "wandb_entity": "robert_jarolim",
-    "wandb_project": "solis",
-    "wandb_name": "solis_401_photospheric"
-  },
-  
-  "data": {
-    "type": "solis",
-    "data_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/multi_height/401.npy",
-    "height_mapping": {
-      "z":  [0.000]
-    },
-    "slices": [0],
-    "Mm_per_pixel": 0.36,
-    "boundary": {"type":  "potential", "strides":  4},
-    "height": 320,
-    "b_norm": 2500,
-    "spatial_norm": 320,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/solis_401_photospheric",
-    "num_workers": 8
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false
-  },
-
-  "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}
-  }
-}
\ No newline at end of file
diff --git a/config/spherical.json b/config/spherical.json
deleted file mode 100644
index 3f60470..0000000
--- a/config/spherical.json
+++ /dev/null
@@ -1,47 +0,0 @@
-{
-  "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
deleted file mode 100644
index 7cd4336..0000000
--- a/config/sst.json
+++ /dev/null
@@ -1,75 +0,0 @@
-{
-  "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": [
-      {
-        "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_v1",
-    "num_workers": 8,
-    "validation_strides": 10
-  },
-  "model": {
-    "dim": 256,
-    "use_height_mapping": true,
-    "use_vector_potential": false
-  },
-  "training": {
-    "epochs": 100,
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "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}
-  }
-}
\ No newline at end of file
diff --git a/config/sst_photospheric.json b/config/sst_photospheric.json
deleted file mode 100644
index dbcc359..0000000
--- a/config/sst_photospheric.json
+++ /dev/null
@@ -1,49 +0,0 @@
-{
-  "base_path": "/gpfs/gpfs0/robert.jarolim/multi_height/sst_photospheric",
-  "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"
-     ],
-    "mask_path": "/gpfs/gpfs0/robert.jarolim/data/nf2/sst/nb_6173_2018-09-30_mosaic_Bvec_cea_mask.fits",
-    "height_mapping": {
-      "z":  [0.000]
-    },
-    "flip_sign": 1,
-    "Mm_per_pixel": 0.36,
-    "boundary": {
-      "type":  "potential", "strides": 4
-    },
-    "b_norm": 2500,
-    "bin": 4,
-    "spatial_norm": 539.5,
-    "height": 539.5,
-    "batch_size": {"boundary":  1e4, "random":  2e4},
-    "iterations": 1e5,
-    "work_directory": "/tmp/nf2/sst_photospheric",
-    "num_workers": 8,
-    "validation_strides": 2
-  },
-
-  "model": {
-    "dim": 256,
-    "use_height_mapping": false,
-    "use_vector_potential": false
-  },
-
-  "training": {
-    "lambda_b": {"start": 1e3, "end": 1, "iterations" : 5e4},
-    "lambda_div": 1e-1,
-    "lambda_ff": 1e-1,
-    "lambda_min_energy": 0,
-    "lambda_height_reg": 1e-3,
-    "validation_interval": 1e4,
-    "lr_params": {"start": 5e-4, "end": 5e-5, "decay_iterations": 1e5}
-  }
-}
\ No newline at end of file
diff --git a/config/sst_vp.json b/config/sst_vp.json
deleted file mode 100644
index 27fd1ed..0000000
--- a/config/sst_vp.json
+++ /dev/null
@@ -1,75 +0,0 @@
-{
-  "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
deleted file mode 100644
index 1346944..0000000
--- a/config/synoptic/synoptic_series.json
+++ /dev/null
@@ -1,47 +0,0 @@
-{
-  "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
deleted file mode 100644
index 01fc600..0000000
--- a/config/synoptic/template.json
+++ /dev/null
@@ -1,65 +0,0 @@
-{
-  "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/13664/init_1.yaml b/config_yaml/13664/init_1.yaml
new file mode 100644
index 0000000..f20dba4
--- /dev/null
+++ b/config_yaml/13664/init_1.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/init_1"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13664_1"
+logging:
+  project: nf2-cartesian
+  name: NOAA 13664
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/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
diff --git a/config_yaml/13664/init_2.yaml b/config_yaml/13664/init_2.yaml
new file mode 100644
index 0000000..0fa5fea
--- /dev/null
+++ b/config_yaml/13664/init_2.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/init_2"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13664_2"
+logging:
+  project: nf2-cartesian
+  name: NOAA 13664
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509_000000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509_000000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509_000000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.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/13664/series_1.yaml b/config_yaml/13664/series_1.yaml
new file mode 100644
index 0000000..a742aee
--- /dev/null
+++ b/config_yaml/13664/series_1.yaml
@@ -0,0 +1,57 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/series"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13664_series"
+meta_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/init_1/last.ckpt"
+logging:
+  project: AR 13664
+  name: series
+data:
+  type: sharp
+  data_path: [
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240505*Bt_err.fits'
+    },
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240506*Bt_err.fits'
+    },
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240507*Bt_err.fits'
+    },
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240508*Bt_err.fits'
+    }
+  ]
+  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/13664/series_2.yaml b/config_yaml/13664/series_2.yaml
new file mode 100644
index 0000000..8e0350b
--- /dev/null
+++ b/config_yaml/13664/series_2.yaml
@@ -0,0 +1,41 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/series_2"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13664_series_2"
+meta_path: "/glade/work/rjarolim/nf2/sharp/13664_v4/init_2/last.ckpt"
+logging:
+  project: AR 13664
+  name: series
+data:
+  type: sharp
+  data_path: [
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240509*Bt_err.fits'
+    },
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.2024051*Bt_err.fits'
+    }
+  ]
+  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/disambiguation/muram_mfr_1slices.yaml b/config_yaml/disambiguation/muram_mfr_1slices.yaml
new file mode 100644
index 0000000..867abcd
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_1slices.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_1slices_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_1slices_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1
+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: True
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
diff --git a/config_yaml/disambiguation/muram_mfr_1slices_ambiguous.yaml b/config_yaml/disambiguation/muram_mfr_1slices_ambiguous.yaml
new file mode 100644
index 0000000..465dd11
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_1slices_ambiguous.yaml
@@ -0,0 +1,41 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_1slices_ambiguous_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_1slices_ambiguous_v02"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1 - ambiguous v2
+data:
+  type: muram
+  slices:
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.100.474000"
+      los_trv_azi_transform: True
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: azimuth_disambiguation
+      lambda: { start: 1.0, end: 1.0e-2, iterations: 1.0e+5 }
+      ds_id: [ boundary_01 ]
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
diff --git a/config_yaml/disambiguation/muram_mfr_2slices.yaml b/config_yaml/disambiguation/muram_mfr_2slices.yaml
new file mode 100644
index 0000000..0423cbb
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_2slices.yaml
@@ -0,0 +1,43 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_2slices_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_2slices_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1 + tau1e-6
+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: True
+      ambiguous_azimuth: False
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 23}
+      los_trv_azi_transform: True
+      ambiguous_azimuth: False
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
diff --git a/config_yaml/disambiguation/muram_mfr_2slices_ambiguous.yaml b/config_yaml/disambiguation/muram_mfr_2slices_ambiguous.yaml
new file mode 100644
index 0000000..7c03977
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_2slices_ambiguous.yaml
@@ -0,0 +1,54 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_2slices_ambiguous_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_2slices_ambiguous_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1 + tau1e-6 - ambiguous
+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: True
+      ambiguous_azimuth: True
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 23}
+      los_trv_azi_transform: True
+      ambiguous_azimuth: True
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/disambiguation/muram_mfr_3slices.yaml b/config_yaml/disambiguation/muram_mfr_3slices.yaml
new file mode 100644
index 0000000..c3b1f33
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_3slices.yaml
@@ -0,0 +1,46 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_3slices_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_3slices_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1 + tau1e-4 + tau1e-6
+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: True
+      ambiguous_azimuth: False
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000100.474000"
+      height_mapping: { z: 1.0, z_min: 0, z_max: 4 }
+      los_trv_azi_transform: True
+      ambiguous_azimuth: False
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      los_trv_azi_transform: True
+      ambiguous_azimuth: False
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02, boundary_03 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 , boundary_03 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02, boundary_03, validation_boundary_03 ]
diff --git a/config_yaml/disambiguation/muram_mfr_3slices_ambiguous.yaml b/config_yaml/disambiguation/muram_mfr_3slices_ambiguous.yaml
new file mode 100644
index 0000000..491d4c1
--- /dev/null
+++ b/config_yaml/disambiguation/muram_mfr_3slices_ambiguous.yaml
@@ -0,0 +1,62 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_mfr_3slices_ambiguous_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_mfr_3slices_ambiguous_v02"
+logging:
+  project: auto-disambiguation
+  name: MURAM MFR - tau1 + tau1e-4 + tau1e-6 - ambiguous
+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: True
+      ambiguous_azimuth: True
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.010.474000"
+      height_mapping: { z: 0.5, z_min: 0, z_max: 10 }
+      los_trv_azi_transform: True
+      scaling: 1.5
+      ambiguous_azimuth: True
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/2D/tau_slice_0.000001.474000"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      los_trv_azi_transform: True
+      ambiguous_azimuth: True
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02, boundary_03 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: azimuth_disambiguation
+      lambda: { start: 1.0, end: 1.0e-2, iterations: 1.0e+5 }
+      ds_id: [ boundary_01, boundary_02, boundary_03 ]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 , boundary_03 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02, boundary_03, validation_boundary_03 ]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 , boundary_03, validation_boundary_03 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
+  - type: disambiguation
+    name: disambiguation-03
+    ds_id: validation_boundary_03
diff --git a/config_yaml/disambiguation/muram_sma_1slices.yaml b/config_yaml/disambiguation/muram_sma_1slices.yaml
new file mode 100644
index 0000000..59a6ddf
--- /dev/null
+++ b/config_yaml/disambiguation/muram_sma_1slices.yaml
@@ -0,0 +1,32 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_sma_1slices_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_sma_1slices_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM SMA - tau1
+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: True
+      ambiguous_azimuth: False
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
diff --git a/config_yaml/disambiguation/muram_sma_2slices.yaml b/config_yaml/disambiguation/muram_sma_2slices.yaml
new file mode 100644
index 0000000..9f9fd42
--- /dev/null
+++ b/config_yaml/disambiguation/muram_sma_2slices.yaml
@@ -0,0 +1,43 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_sma_2slices_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_sma_2slices_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM SMA - tau1 + tau1e-6
+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: True
+      ambiguous_azimuth: False
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_0.000001.399000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 23}
+      los_trv_azi_transform: True
+      ambiguous_azimuth: False
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: height
+      name: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
diff --git a/config_yaml/disambiguation/muram_sma_2slices_ambiguous.yaml b/config_yaml/disambiguation/muram_sma_2slices_ambiguous.yaml
new file mode 100644
index 0000000..b47111c
--- /dev/null
+++ b/config_yaml/disambiguation/muram_sma_2slices_ambiguous.yaml
@@ -0,0 +1,54 @@
+---
+base_path: "/glade/work/rjarolim/nf2/disambiguation/muram_sma_2slices_ambiguous_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/muram_sma_2slices_ambiguous_v01"
+logging:
+  project: auto-disambiguation
+  name: MURAM SMA - tau1 + tau1e-6 - ambiguous
+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: True
+      ambiguous_azimuth: True
+    - data_path: "/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/2D/tau_slice_0.000001.399000"
+      height_mapping: {z: 2.0, z_min: 0, z_max: 23}
+      los_trv_azi_transform: True
+      ambiguous_azimuth: True
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      name: boundary
+      lambda: 1.0
+      ds_id: [ boundary_01, boundary_02 ]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/multi_height/muram_mfr_disambiguation.yaml b/config_yaml/multi_height/muram_mfr_disambiguation.yaml
deleted file mode 100644
index 8600bad..0000000
--- a/config_yaml/multi_height/muram_mfr_disambiguation.yaml
+++ /dev/null
@@ -1,58 +0,0 @@
----
-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/sharp/13229.yaml b/config_yaml/sharp/13229.yaml
index a306602..4dec67d 100644
--- a/config_yaml/sharp/13229.yaml
+++ b/config_yaml/sharp/13229.yaml
@@ -1,25 +1,25 @@
 ---
-base_path: "/glade/work/rjarolim/nf2/sharp/13229_custom_v02"
-work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13229_v01"
+base_path: "/glade/work/rjarolim/nf2/sharp/13229/init"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13229_init"
 logging:
-  project: nf2-cartesian
-  name: NOAA 13229
+  project: AR 13229
+  name: init
 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"
+        Br: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_IN.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_IN.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_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"
+        Br_err: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_IN.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_IN.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13229_large/hmi.Bptr_720s_mod.20230224_100000_TAI.JSOC_20240815_000809_IN.Bp_err.fits"
   num_workers: 8
   iterations: 10000
-  max_height: 150
+  z_range: [0, 200]
 model:
-  type: vector_potential
+  type: b
   dim: 256
 training:
   epochs: 15
@@ -30,3 +30,5 @@ training:
       ds_id: [boundary_01, potential]
     - type: force_free
       lambda: 1.0e-1
+    - type: divergence
+      lambda: 1.0e-1
diff --git a/config_yaml/sharp/13229_series.yaml b/config_yaml/sharp/13229_series.yaml
new file mode 100644
index 0000000..31d859f
--- /dev/null
+++ b/config_yaml/sharp/13229_series.yaml
@@ -0,0 +1,36 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13229/series"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13229_series"
+meta_path: "/glade/work/rjarolim/nf2/sharp/13229/init/last.ckpt"
+logging:
+  project: AR 13229
+  name: series
+data:
+  type: sharp
+  data_path: [
+    {
+      Bp: '/glade/work/rjarolim/data/nf2/13229_large/*Bp.fits',
+      Br: '/glade/work/rjarolim/data/nf2/13229_large/*Br.fits',
+      Bt: '/glade/work/rjarolim/data/nf2/13229_large/*Bt.fits',
+      Bp_err: '/glade/work/rjarolim/data/nf2/13229_large/*Bp_err.fits',
+      Br_err: '/glade/work/rjarolim/data/nf2/13229_large/*Br_err.fits',
+      Bt_err: '/glade/work/rjarolim/data/nf2/13229_large/*Bt_err.fits'
+    }
+  ]
+  num_workers: 8
+  iterations: 5.0e+3
+  z_range: [0, 200]
+model:
+  type: b
+  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
+    - type: divergence
+      lambda: 1.0e-1
\ No newline at end of file
diff --git a/config_yaml/sharp/13664.yaml b/config_yaml/sharp/13664.yaml
deleted file mode 100644
index ccad7c1..0000000
--- a/config_yaml/sharp/13664.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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/ensemble_20240510_0600.yaml b/config_yaml/sharp/13664_ensemble/ensemble_20240510_0600.yaml
new file mode 100644
index 0000000..c0c71b4
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/ensemble_20240510_0600.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_ensemble/20240510_0600_{run_id}"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13664_ensemble/20240510_0600_{run_id}"
+logging:
+  project: Ensemble 13664
+  name: "20240510_0600_{run_id}"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240510_060000_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/ensemble_20240511_0100.yaml b/config_yaml/sharp/13664_ensemble/ensemble_20240511_0100.yaml
new file mode 100644
index 0000000..37b56ad
--- /dev/null
+++ b/config_yaml/sharp/13664_ensemble/ensemble_20240511_0100.yaml
@@ -0,0 +1,31 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/13664_ensemble/20240511_0100_{run_id}"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13664_ensemble/20240511_0100_{run_id}"
+logging:
+  project: Ensemble 13664
+  name: "20240511_0100_{run_id}"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_TAI.Bp.fits"
+      error_path:
+        Br_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/nf2/13664/hmi.sharp_cea_720s.11149.20240511_010000_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/run01.yaml b/config_yaml/sharp/13664_ensemble/run01.yaml
deleted file mode 100644
index 720ca58..0000000
--- a/config_yaml/sharp/13664_ensemble/run01.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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
deleted file mode 100644
index 511c1d9..0000000
--- a/config_yaml/sharp/13664_ensemble/run02.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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
deleted file mode 100644
index 8c5f6e1..0000000
--- a/config_yaml/sharp/13664_ensemble/run03.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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
deleted file mode 100644
index 36f4de5..0000000
--- a/config_yaml/sharp/13664_ensemble/run04.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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
deleted file mode 100644
index d918e2f..0000000
--- a/config_yaml/sharp/13664_ensemble/run05.yaml
+++ /dev/null
@@ -1,31 +0,0 @@
----
-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
deleted file mode 100644
index b34b943..0000000
--- a/config_yaml/sharp/13664_series.yaml
+++ /dev/null
@@ -1,24 +0,0 @@
----
-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_magnetostatic.yaml b/config_yaml/sharp/377_magnetostatic.yaml
new file mode 100644
index 0000000..e69de29
diff --git a/config_yaml/sharp/377_magnetostatic_v2.yaml b/config_yaml/sharp/377_magnetostatic_v2.yaml
new file mode 100644
index 0000000..bc0ad02
--- /dev/null
+++ b/config_yaml/sharp/377_magnetostatic_v2.yaml
@@ -0,0 +1,36 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sharp/377_magnetostatic_v03"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_magnetostatic_v03"
+logging:
+  project: nf2-cartesian
+  name: magnetostatic - no divergence regularization
+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"
+      batch_size: 2048
+  random_config:
+    batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+  batch_size: 1024
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [boundary_01, potential]
+    - type: implicit_magnetostatic
+      lambda: 1.0e-3
+      divergence_regularization: False
diff --git a/config_yaml/sharp/377_series.yaml b/config_yaml/sharp/377_series.yaml
index 6389dd1..b0f93c7 100644
--- a/config_yaml/sharp/377_series.yaml
+++ b/config_yaml/sharp/377_series.yaml
@@ -11,9 +11,8 @@ data:
   num_workers: 8
   iterations: 2.0e+3
 model:
-  type: vector_potential
+  type: b
   dim: 256
-  activation: swish
 training:
   epochs: 1.0e+3
   check_val_every_n_epoch: 5 # validation plots in 1h steps
@@ -23,4 +22,6 @@ training:
       lambda: 1
       ds_id: [boundary, potential]
     - type: force_free
+      lambda: 1.0e-1
+    - type: divergence
       lambda: 1.0e-1
\ No newline at end of file
diff --git a/config_yaml/spherical/2173.yaml b/config_yaml/spherical/2173.yaml
index dabc768..251f02c 100644
--- a/config_yaml/spherical/2173.yaml
+++ b/config_yaml/spherical/2173.yaml
@@ -1,9 +1,9 @@
 ---
-base_path: "/glade/work/rjarolim/nf2/spherical/2173_v03"
-work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_v03"
+base_path: "/glade/work/rjarolim/nf2/spherical/2173/{time}_v05"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/{time}_v05"
 logging:
   project: spherical
-  name: "2173"
+  name: "2173 {time}"
 data:
   type: spherical
   train_configs:
@@ -11,22 +11,35 @@ data:
       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"
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bp.fits"
+#       Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Br_err.fits"
+#       Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bt_err.fits"
+#       Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bp_err.fits"
+#      mask_configs:
+#        type: heliographic_carrington
+#        longitude_range: [ 1.446, 4.273 ]
+#        latitude_range: [ 0.424, 2.938 ]
+#        invert: true
+    - 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: heliographic_carrington
-        longitude_range: [ 1.446, 4.273 ]
-        latitude_range: [  0.424, 2.938 ]
-        invert: true
+#        type: heliographic_carrington
+#        longitude_range: [ 1.446, 4.273 ]
+#        latitude_range: [ 0.424, 2.938 ]
+        type: reference
+        file: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Br.fits"
     - type: random_spherical
       ds_id: random
-      batch_size: 8192
-      longitude_range: [ 1.446, 4.273 ]
-      latitude_range: [  0.424, 2.938 ]
+      batch_size: 16384
+#      longitude_range: [ 1.446, 4.273 ]
+#      latitude_range: [  0.424, 2.938 ]
   validation_configs:
     - type: map
       ds_id: full_disk_valid
@@ -34,16 +47,16 @@ data:
       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"
+        Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.{time}_TAI.Bp_err.fits"
       mask_configs:
         type: heliographic_carrington
         longitude_range: [ 1.446, 4.273 ]
-        latitude_range: [  0.424, 2.938 ]
+        latitude_range: [ 0.424, 2.938 ]
         invert: true
     - type: spherical_slices
       ds_id: slices
@@ -64,14 +77,14 @@ model:
   type: vector_potential
   dim: 512
 training:
-  epochs: 100
+  epochs: 1000
   loss_config:
     - type: boundary
       name: boundary
-      lambda: 1 #{start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
-      ds_id: full_disk
+      lambda: {start: 1.0e+3, end: 20.0, iterations: 5.0e+4}
+      ds_id: [full_disk, synoptic]
     - type: energy_gradient
-      lambda: 1.0
+      lambda: 10.0
       ds_id: random
       base_radius: 1.04
     - type: potential
diff --git a/config_yaml/spherical/2173_free.yaml b/config_yaml/spherical/2173_free.yaml
new file mode 100644
index 0000000..52e4e09
--- /dev/null
+++ b/config_yaml/spherical/2173_free.yaml
@@ -0,0 +1,87 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_free_v06"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_free_v06"
+logging:
+  project: spherical
+  name: "2173 free"
+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_000000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_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: 256
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4}
+      ds_id: [full_disk, synoptic]
+    - type: energy_gradient
+      lambda: 0.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 0.0
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
+    - type: divergence
+      lambda: 0.0
diff --git a/config_yaml/spherical/2173_full.yaml b/config_yaml/spherical/2173_full.yaml
index a3bbfb0..91a4e75 100644
--- a/config_yaml/spherical/2173_full.yaml
+++ b/config_yaml/spherical/2173_full.yaml
@@ -1,6 +1,6 @@
 ---
-base_path: "/glade/work/rjarolim/nf2/spherical/2173_full_divergence"
-work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_full_divergence"
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_full_15_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_full_15_v02"
 logging:
   project: spherical
   name: "2173 full"
@@ -11,9 +11,9 @@ data:
       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: "/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: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_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"
@@ -26,7 +26,7 @@ data:
         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"
+        file: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits"
     - type: random_spherical
       ds_id: random
       batch_size: 16384
@@ -37,15 +37,15 @@ data:
       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"
+        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: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp.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"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp_err.fits"
       mask_configs:
         type: heliographic_carrington
-        longitude_range: [ 1.446, 4.273 ]
+        longitude_range: [ 1.603, 4.116 ]
         latitude_range: [  0.424, 2.938 ]
         invert: true
     - type: spherical_slices
@@ -54,27 +54,32 @@ data:
       plot_currents: true
       longitude_range: [ 1.446, 4.273 ]
       latitude_range: [  0.424, 2.938 ]
+    - type: spherical_slices
+      ds_id: synoptic_slices
+      n_slices: 10
+      plot_currents: true
     - type: sphere
       ds_id: sphere
       longitude_range: [ 1.446, 4.273 ]
       latitude_range: [  0.424, 2.938 ]
-  max_radius: 1.3
+  max_radius: 1.5
   G_per_dB: 2000
   num_workers: 8
   batch_size: 16384
-  iterations: 10000
+  iterations: 5000
 model:
   type: vector_potential
+#  activation: swish
   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}
+      lambda: {start: 1.0e+3, end: 10.0, iterations: 5.0e+4}
       ds_id: [full_disk, synoptic]
     - type: energy_gradient
-      lambda: 1.0
+      lambda: 10.0
       ds_id: random
       base_radius: 1.04
     - type: potential
@@ -83,3 +88,5 @@ training:
       base_radius: 1.04
     - type: force_free
       lambda: 0.1
+    - type: divergence
+      lambda: 0.0
diff --git a/config_yaml/spherical/2173_insert.yaml b/config_yaml/spherical/2173_insert.yaml
new file mode 100644
index 0000000..c4af058
--- /dev/null
+++ b/config_yaml/spherical/2173_insert.yaml
@@ -0,0 +1,72 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_insert_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_insert_v02"
+logging:
+  project: spherical
+  name: "2173 insert"
+data:
+  type: spherical
+  train_configs:
+    - 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"
+      insert:
+        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"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+  validation_configs:
+    - type: map
+      ds_id: synoptic_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"
+      insert:
+        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"
+    - 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: 256
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 1
+      ds_id: synoptic
+    - type: energy_gradient
+      lambda: 10.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_post.yaml b/config_yaml/spherical/2173_post.yaml
new file mode 100644
index 0000000..71c18e3
--- /dev/null
+++ b/config_yaml/spherical/2173_post.yaml
@@ -0,0 +1,92 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_post_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_post_v02"
+logging:
+  project: spherical
+  name: "2173 post"
+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_230000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_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_230000_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_230000_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bp.fits"
+        Br_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Br_err.fits"
+        Bt_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bt_err.fits"
+        Bp_err: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bp_err.fits"
+      mask_configs:
+        type: heliographic_carrington
+        longitude_range: [ 1.603, 4.116 ]
+        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: spherical_slices
+      ds_id: synoptic_slices
+      n_slices: 10
+      plot_currents: true
+    - type: sphere
+      ds_id: sphere
+      longitude_range: [ 1.446, 4.273 ]
+      latitude_range: [  0.424, 2.938 ]
+  max_radius: 1.5
+  G_per_dB: 2000
+  num_workers: 8
+  batch_size: 16384
+  iterations: 5000
+model:
+  type: vector_potential
+#  activation: swish
+  dim: 512
+training:
+  epochs: 1000
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: {start: 1.0e+3, end: 10.0, iterations: 5.0e+4}
+      ds_id: [full_disk, synoptic]
+    - type: energy_gradient
+      lambda: 10.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 0.0
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
+    - type: divergence
+      lambda: 0.0
diff --git a/config_yaml/spherical/2173_series.yaml b/config_yaml/spherical/2173_series.yaml
new file mode 100644
index 0000000..0aeb184
--- /dev/null
+++ b/config_yaml/spherical/2173_series.yaml
@@ -0,0 +1,163 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_series_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_series_v02"
+meta_path: "/glade/work/rjarolim/nf2/spherical/2173/20160205_160000_v04/last.ckpt"
+logging:
+  project: spherical
+  name: "2173 series"
+data:
+  type: spherical
+  data_path:
+    - 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: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_171200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_171200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_171200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_172400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_172400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_172400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_173600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_173600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_173600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_174800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_174800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_174800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_180000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_180000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_180000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_181200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_181200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_181200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_182400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_182400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_182400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_183600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_183600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_183600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_184800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_184800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_184800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_190000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_190000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_190000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_191200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_191200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_191200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_192400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_192400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_192400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_193600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_193600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_193600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_194800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_194800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_194800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_200000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_200000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_200000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_201200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_201200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_201200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_202400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_202400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_202400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_203600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_203600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_203600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_204800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_204800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_204800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_210000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_210000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_210000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_211200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_211200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_211200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_212400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_212400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_212400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_213600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_213600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_213600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_214800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_214800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_214800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_220000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_220000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_220000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_221200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_221200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_221200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_222400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_222400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_222400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_223600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_223600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_223600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_224800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_224800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_224800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_230000_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_231200_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_231200_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_231200_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_232400_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_232400_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_232400_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_233600_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_233600_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_233600_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_234800_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_234800_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_234800_TAI.Bp.fits"
+    - Br: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160206_000000_TAI.Br.fits"
+      Bt: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160206_000000_TAI.Bt.fits"
+      Bp: "/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160206_000000_TAI.Bp.fits"
+  synoptic_fits_path:
+      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"
+  validation_configs:
+    - 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: spherical_slices
+      ds_id: full_sphere
+      n_slices: 10
+      plot_currents: true
+    - 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: 2000
+model:
+  type: vector_potential
+  dim: 512
+training:
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 100.0
+      ds_id: [full_disk, synoptic]
+    - type: energy_gradient
+      lambda: 10.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
+    - type: divergence
+      lambda: 0.0
diff --git a/config_yaml/spherical/2173_synoptic.yaml b/config_yaml/spherical/2173_synoptic.yaml
new file mode 100644
index 0000000..a6d8d12
--- /dev/null
+++ b/config_yaml/spherical/2173_synoptic.yaml
@@ -0,0 +1,64 @@
+---
+base_path: "/glade/work/rjarolim/nf2/spherical/2173_synoptic_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/2173_synoptic_v01"
+logging:
+  project: spherical
+  name: "2173 synoptic"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: full_disk
+      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"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 8192
+  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/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: 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: 256
+training:
+  epochs: 100
+  loss_config:
+    - type: boundary
+      name: boundary
+      lambda: 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
\ No newline at end of file
diff --git a/config_yaml/spherical/377_spherical.yaml b/config_yaml/spherical/377_spherical.yaml
index b0ea1b9..c6c4509 100644
--- a/config_yaml/spherical/377_spherical.yaml
+++ b/config_yaml/spherical/377_spherical.yaml
@@ -1,6 +1,6 @@
 ---
-base_path: "/glade/work/rjarolim/nf2/spherical/377_v05"
-work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_v05"
+base_path: "/glade/work/rjarolim/nf2/spherical/comparison/377_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/377_v01"
 logging:
   project: spherical
   name: '377 spherical'
@@ -76,7 +76,7 @@ data:
   G_per_dB: 2000
   num_workers: 8
   batch_size: 16384
-  iterations: 1000
+  iterations: 10000
 model:
   type: vector_potential
   dim: 512
@@ -88,7 +88,7 @@ training:
       lambda: {start: 1.0e+3, end: 1.0, iterations: 5.0e+4} #1
       ds_id: full_disk
     - type: energy_gradient
-      lambda: 1.0
+      lambda: 10.0
       ds_id: random
       base_radius: 1.04
     - type: potential
diff --git a/config_yaml/spherical/synoptic.yaml b/config_yaml/spherical/synoptic.yaml
new file mode 100644
index 0000000..2916712
--- /dev/null
+++ b/config_yaml/spherical/synoptic.yaml
@@ -0,0 +1,64 @@
+---
+base_path: "/glade/work/rjarolim/nf2/synoptic/{CR}_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/{CR}_v01"
+logging:
+  project: synoptic
+  name: "{CR}"
+data:
+  type: spherical
+  train_configs:
+    - type: map
+      ds_id: synoptic
+      batch_size: 8192
+      files:
+        Br: "/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.{CR}.Mr_polfil.fits"
+        Bt: "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{CR}.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{CR}.Bp.fits"
+    - type: random_spherical
+      ds_id: random
+      batch_size: 16384
+  validation_configs:
+    - type: map
+      ds_id: synoptic_valid
+      filter_nans: false
+      shuffle: false
+      plot_overview: false
+      files:
+        Br: "/glade/work/rjarolim/data/global/synoptic/hmi.synoptic_mr_polfil_720s.{CR}.Mr_polfil.fits"
+        Bt: "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{CR}.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/global/synoptic/hmi.b_synoptic.{CR}.Bp.fits"
+    - type: spherical_slices
+      ds_id: slices
+      n_slices: 10
+      plot_currents: true
+    - type: sphere
+      ds_id: sphere
+  max_radius: 2.5
+  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: {start: 1.0e+3, end: 10.0, iterations: 5.0e+4}
+      ds_id: synoptic
+    - type: energy_gradient
+      lambda: 10.0
+      ds_id: random
+      base_radius: 1.04
+    - type: potential
+      lambda: 0.0 #1.0e-2
+      ds_id: random
+      base_radius: 1.04
+    - type: force_free
+      lambda: 0.1
+    - type: radial
+      lambda: 1.0e-4
+      ds_id: random
+      base_radius: 2.0
\ No newline at end of file
diff --git a/config_yaml/sst/13392_1slices_0851.yaml b/config_yaml/sst/13392_1slices_0851.yaml
new file mode 100644
index 0000000..9f576e7
--- /dev/null
+++ b/config_yaml/sst/13392_1slices_0851.yaml
@@ -0,0 +1,50 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_1slices_0851_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_1slices_0851_v02"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - 6173"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 8192
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 0.0 #1.0e-4
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01 ]
+  check_val_every_n_epoch: 1
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01 ]
+callbacks:
+  - type: disambiguation
+    ds_id: validation_boundary_01
diff --git a/config_yaml/sst/13392_2slices7699_0851.yaml b/config_yaml/sst/13392_2slices7699_0851.yaml
new file mode 100644
index 0000000..d4150e9
--- /dev/null
+++ b/config_yaml/sst/13392_2slices7699_0851.yaml
@@ -0,0 +1,69 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_2slices7699_0851_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_2slices7699_0851_v02"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - 6173+7699"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_7699_StkIQUV_rebin_v2.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_7699_StkIQUV_rebin_v2.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_7699_StkIQUV_rebin_v2.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: {z: 0.5, z_min: 0, z_max: 10}
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/sst/13392_2slices_0851.yaml b/config_yaml/sst/13392_2slices_0851.yaml
new file mode 100644
index 0000000..67712ab
--- /dev/null
+++ b/config_yaml/sst/13392_2slices_0851.yaml
@@ -0,0 +1,69 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_2slices_0851_v03"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_2slices_0851_v03"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - 6173+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 0.0 # 1.0e-4
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/sst/13392_2slices_1050.yaml b/config_yaml/sst/13392_2slices_1050.yaml
new file mode 100644
index 0000000..0f8a4ea
--- /dev/null
+++ b/config_yaml/sst/13392_2slices_1050.yaml
@@ -0,0 +1,69 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_2slices_1050_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_2slices_1050_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SST 10:50 - 6173+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BLOS_panorama1050_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BTRV_panorama1050_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BAZI_panorama1050_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BLOS_panorama1050_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BTRV_panorama1050_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/BAZI_panorama1050_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_1050_converted/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/sst/13392_2slices_7699_8542_0851.yaml b/config_yaml/sst/13392_2slices_7699_8542_0851.yaml
new file mode 100644
index 0000000..46f464a
--- /dev/null
+++ b/config_yaml/sst/13392_2slices_7699_8542_0851.yaml
@@ -0,0 +1,69 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_2slices_7699_8542_0851_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_2slices_7699_8542_0851_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - - 7699+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+          B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_7699_StkIQUV_rebin_v2.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02 ]
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, validation_boundary_01, validation_boundary_02 ]
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/sst/13392_3slicesLOS_0851.yaml b/config_yaml/sst/13392_3slicesLOS_0851.yaml
new file mode 100644
index 0000000..bafee7c
--- /dev/null
+++ b/config_yaml/sst/13392_3slicesLOS_0851.yaml
@@ -0,0 +1,84 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_3slicesLOS_0851_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_3slicesLOS_0851_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - - 6173+5896+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_5896_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_5896_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_5896_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 1.0, z_min: 0, z_max: 15 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_03]
+    - type: boundary_los
+      lambda: 1.0
+      ds_id: [boundary_02]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_03 ]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02, boundary_03 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02,
+               boundary_03, validation_boundary_03 ]
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01,
+               boundary_02, validation_boundary_02,
+               boundary_03, validation_boundary_03 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
\ No newline at end of file
diff --git a/config_yaml/sst/13392_3slices_0851.yaml b/config_yaml/sst/13392_3slices_0851.yaml
new file mode 100644
index 0000000..fa0a21a
--- /dev/null
+++ b/config_yaml/sst/13392_3slices_0851.yaml
@@ -0,0 +1,83 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_3slices_0851_v04"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_3slices_0851_v04"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - - 6173+7699+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+          B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_7699_StkIQUV_rebin_v2.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: {z: 0.5, z_min: 0, z_max: 10}
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02, boundary_03]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02, boundary_03 ]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02, boundary_03 ]
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, boundary_02, boundary_03, validation_boundary_01, validation_boundary_02, validation_boundary_03 ]
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02, boundary_03, validation_boundary_03 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
+  - type: disambiguation
+    name: disambiguation-03
+    ds_id: validation_boundary_03
\ No newline at end of file
diff --git a/config_yaml/sst/13392_4slices_0851.yaml b/config_yaml/sst/13392_4slices_0851.yaml
new file mode 100644
index 0000000..1569cd0
--- /dev/null
+++ b/config_yaml/sst/13392_4slices_0851.yaml
@@ -0,0 +1,102 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_4slices_0851_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_4slices_0851_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - - 6173+7699+5896+8542"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_6173_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_6173_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_6173_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+          B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_7699_StkIQUV_rebin_v2.fits"
+          B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_7699_StkIQUV_rebin_v2.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: {z: 0.5, z_min: 0, z_max: 10}
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_5896_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_5896_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_5896_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 1.0, z_min: 0, z_max: 15 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_8542_StkIQUV_rebin.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_8542_StkIQUV_rebin.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_8542_StkIQUV_rebin.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      height_mapping: { z: 2.0, z_min: 0, z_max: 23 }
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 4096
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01, boundary_02, boundary_04]
+    - type: boundary_los
+      lambda: 1.0
+      ds_id: [boundary_03]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01, boundary_02, boundary_04 ]
+    - type: min_height
+      lambda: 1.0e-3
+      ds_id: [ boundary_02, boundary_03, boundary_04 ]
+  transforms:
+    - type: height
+      ds_id: [ boundary_02, validation_boundary_02,
+               boundary_03, validation_boundary_03,
+               boundary_04, validation_boundary_04 ]
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01,
+               boundary_02, validation_boundary_02,
+               boundary_03, validation_boundary_03,
+               boundary_04, validation_boundary_04 ]
+callbacks:
+  - type: disambiguation
+    name: disambiguation-01
+    ds_id: validation_boundary_01
+  - type: disambiguation
+    name: disambiguation-02
+    ds_id: validation_boundary_02
+  - type: disambiguation
+    name: disambiguation-04
+    ds_id: validation_boundary_04
\ No newline at end of file
diff --git a/config_yaml/sst/13392_7699_0851.yaml b/config_yaml/sst/13392_7699_0851.yaml
new file mode 100644
index 0000000..1a4009d
--- /dev/null
+++ b/config_yaml/sst/13392_7699_0851.yaml
@@ -0,0 +1,50 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/13392_7699_0851_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/13392_7699_0851_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SST 08:51 - 7699"
+data:
+  type: fits
+  slices:
+    - type: los_trv_azi
+      Mm_per_pixel: 0.09
+      fits_path:
+        B_los: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_7699_StkIQUV_rebin_v2.fits"
+        B_trv: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_7699_StkIQUV_rebin_v2.fits"
+        B_azi: "/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_7699_StkIQUV_rebin_v2.fits"
+      mask_path: "/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits"
+      load_map: False
+      ambiguous_azimuth: True
+      batch_size: 8192
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: 1.0e-3
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01 ]
+  check_val_every_n_epoch: 1
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01 ]
+callbacks:
+  - type: disambiguation
+    ds_id: validation_boundary_01
diff --git a/config_yaml/sst/sharp_13392.yaml b/config_yaml/sst/sharp_13392.yaml
new file mode 100644
index 0000000..3c698ff
--- /dev/null
+++ b/config_yaml/sst/sharp_13392.yaml
@@ -0,0 +1,34 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/sharp_13392_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13392_v02"
+logging:
+  project: nf2-sst
+  name: "13392 SHARP 09:12"
+data:
+  type: fits
+  slices:
+    - type: fld_inc_azi
+      fits_path:
+        B_fld: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.field.fits"
+        B_inc: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.inclination.fits"
+        B_azi: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.azimuth.fits"
+#        B_amb: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.disambig.fits"
+      batch_size: 8192
+      slice: [380, 840, 0, 1000]
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: boundary_01
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
diff --git a/config_yaml/sst/sharp_13392_ambiguous.yaml b/config_yaml/sst/sharp_13392_ambiguous.yaml
new file mode 100644
index 0000000..1e6d715
--- /dev/null
+++ b/config_yaml/sst/sharp_13392_ambiguous.yaml
@@ -0,0 +1,49 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/sharp_13392_ambiguous_v02"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13392_ambiguous_v02"
+logging:
+  project: nf2-sst
+  name: "13392 SHARP 09:12 - ambiguous"
+data:
+  type: fits
+  slices:
+    - type: fld_inc_azi
+      fits_path:
+        B_fld: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.field.fits"
+        B_inc: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.inclination.fits"
+        B_azi: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.azimuth.fits"
+#        B_amb: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_720s.9875.20230806_091200_TAI.disambig.fits"
+      ambiguous_azimuth: True
+      batch_size: 8192
+      slice: [380, 840, 0, 1000]
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary_los_trv_azi
+      lambda: 1.0
+      ds_id: [boundary_01]
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
+    - type: potential
+      lambda: {start: 1.0e-2, end: 1.0e-4, iterations: 1.0e+5}
+    - type: divergence
+      lambda:  0.0
+    - type: azimuth_disambiguation
+      lambda: {start: 1.0, end: 1.0e-2, iterations: 1.0e+5}
+      ds_id: [ boundary_01 ]
+  check_val_every_n_epoch: 1
+  transforms:
+    - type: azimuth
+      ds_id: [ boundary_01, validation_boundary_01 ]
+callbacks:
+  - type: disambiguation
+    ds_id: validation_boundary_01
diff --git a/config_yaml/sst/sharp_13392_cea.yaml b/config_yaml/sst/sharp_13392_cea.yaml
new file mode 100644
index 0000000..4d7fd69
--- /dev/null
+++ b/config_yaml/sst/sharp_13392_cea.yaml
@@ -0,0 +1,32 @@
+---
+base_path: "/glade/work/rjarolim/nf2/sst/sharp_13392_cea_v01"
+work_directory: "/glade/derecho/scratch/rjarolim/nf2/sharp_13392_cea_v01"
+logging:
+  project: nf2-sst
+  name: "13392 SHARP CEA 09:12"
+data:
+  type: fits
+  slices:
+    - fits_path:
+        Br: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_cea_720s.9875.20230806_091200_TAI.Br.fits"
+        Bt: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_cea_720s.9875.20230806_091200_TAI.Bt.fits"
+        Bp: "/glade/work/rjarolim/data/nf2/13392/hmi.sharp_cea_720s.9875.20230806_091200_TAI.Bp.fits"
+      batch_size: 8192
+      slice: [430, 900, 0, 1000]
+  num_workers: 8
+  iterations: 10000
+model:
+  type: vector_potential
+  dim: 256
+training:
+  epochs: 15
+  loss_config:
+    - type: boundary
+      lambda: 1.0
+      ds_id: boundary_01
+    - type: boundary
+      name: potential_boundary
+      lambda: 0.1
+      ds_id: potential
+    - type: force_free
+      lambda: 1.0e-2
diff --git a/nf2/convert/convert_series.py b/nf2/convert/convert_series.py
index 558da89..1856b71 100644
--- a/nf2/convert/convert_series.py
+++ b/nf2/convert/convert_series.py
@@ -7,23 +7,28 @@
 
 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('--nf2_dir', type=str, help='path to the source NF2 files', nargs='+', required=True)
     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('--metrics', type=str, nargs='*', help='metrics to be computed', required=False, default=['j'])
     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))
+
+    nf2_dir = args.nf2_dir if isinstance(args.nf2_dir, list) else [args.nf2_dir]
+    nf2_paths = [sorted(glob.glob(f)) for f in nf2_dir]
+    nf2_paths = [f for files in nf2_paths for f in files]  # flatten list
 
     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)
+    out_dir = args.out_dir if args.out_dir is not None else os.path.dirname(nf2_dir[0])
     os.makedirs(out_dir, exist_ok=True)
     height_range = args.height_range
+    metrics = args.metrics
 
     conversion_type = args.type
     if conversion_type == 'vtk':
@@ -51,7 +56,8 @@ def main():
             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)
+                  Mm_per_pixel=Mm_per_pixel, height_range=height_range, metrics=metrics,
+                  progress=False)
 
 
 if __name__ == '__main__':
diff --git a/nf2/convert/muram_to_vtk.py b/nf2/convert/muram_to_vtk.py
new file mode 100644
index 0000000..373d574
--- /dev/null
+++ b/nf2/convert/muram_to_vtk.py
@@ -0,0 +1,33 @@
+import argparse
+
+import numpy as np
+from astropy import constants as const
+from astropy import units as u
+
+from nf2.evaluation.metric import curl
+from nf2.evaluation.vtk import save_vtk
+from nf2.loader.muram import MURaMSnapshot
+
+
+def main():
+    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 * 4)
+    parser.add_argument('--height', type=float, help='height of the snapshot', required=False, default=None)
+    args = parser.parse_args()
+
+    snapshot = MURaMSnapshot(args.source_path, args.iteration)
+    muram_cube = snapshot.load_cube(args.Mm_per_pixel * u.Mm / u.pix, target_tau=1.0, height=args.height * u.Mm)
+    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)
+
+    save_vtk(args.vtk_path, vectors={'b': b, 'j': j.value},
+             scalars={'tau': muram_cube['tau']}, Mm_per_pix=args.Mm_per_pixel)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/nf2/convert/nf2_to_fits.py b/nf2/convert/nf2_to_fits.py
index 7abc4f3..77659d6 100644
--- a/nf2/convert/nf2_to_fits.py
+++ b/nf2/convert/nf2_to_fits.py
@@ -3,7 +3,7 @@
 
 from astropy.io import fits
 
-from nf2.evaluation.output import CartesianOutput, current_density
+from nf2.evaluation.output import CartesianOutput
 
 
 def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
diff --git a/nf2/convert/nf2_to_hdf5.py b/nf2/convert/nf2_to_hdf5.py
index ca80824..35b4ace 100644
--- a/nf2/convert/nf2_to_hdf5.py
+++ b/nf2/convert/nf2_to_hdf5.py
@@ -4,7 +4,7 @@
 
 import h5py
 
-from nf2.evaluation.output import CartesianOutput, current_density
+from nf2.evaluation.output import CartesianOutput
 
 
 class _SaveFileTask(Thread):
@@ -14,6 +14,7 @@ def __init__(self, out_path, output, nf2_out):
         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')
@@ -39,7 +40,6 @@ def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwa
     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')
diff --git a/nf2/convert/nf2_to_vtk.py b/nf2/convert/nf2_to_vtk.py
index 28cf17a..c17a715 100644
--- a/nf2/convert/nf2_to_vtk.py
+++ b/nf2/convert/nf2_to_vtk.py
@@ -8,25 +8,33 @@
 
 class _SaveFileTask(Thread):
 
-    def __init__(self, out_path, output):
+    def __init__(self, out_path, output, metrics):
         super().__init__()
         self.out_path = out_path
         self.output = output
+        self.metrics = metrics if metrics is not None else []
 
     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'])
+        Mm_per_pixel = self.output['Mm_per_pixel']
 
+        # split output into vectors and scalars
+        vectors = {k: v for k, v in self.output['metrics'].items() if len(v.shape) == 4 and v.shape[-1] == 3}
+        scalars = {k: v for k, v in self.output['metrics'].items() if len(v.shape) == 3}
 
-def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, **kwargs):
+        vectors['b'] = self.output['b']
+
+        save_vtk(self.out_path, vectors=vectors, scalars=scalars, Mm_per_pix=Mm_per_pixel)
+
+
+def convert(nf2_path, out_path=None, Mm_per_pixel=None, height_range=None, metrics=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)
+    output = nf2_out.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range, metrics=metrics, **kwargs)
 
     # save file in background
-    task = _SaveFileTask(out_path, output)
+    task = _SaveFileTask(out_path, output, metrics)
     task.start()
 
 
@@ -37,6 +45,7 @@ def main():
     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('--metrics', type=str, nargs='*', help='metrics to be computed', required=False, default=['j'])
 
     args = parser.parse_args()
     nf2_path = args.nf2_path
@@ -44,8 +53,9 @@ def main():
     Mm_per_pixel = args.Mm_per_pixel
     out_path = args.out_path
     height_range = args.height_range
+    metrics = args.metrics
 
-    convert(nf2_path, out_path, Mm_per_pixel, height_range)
+    convert(nf2_path, out_path, Mm_per_pixel, height_range, metrics=metrics, progress=True)
 
 
 if __name__ == '__main__':
diff --git a/nf2/convert/nf2_to_vtk_spherical.py b/nf2/convert/nf2_to_vtk_spherical.py
index 6d1c5e4..47bce67 100644
--- a/nf2/convert/nf2_to_vtk_spherical.py
+++ b/nf2/convert/nf2_to_vtk_spherical.py
@@ -1,72 +1,51 @@
 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.output import SphericalOutput
+from nf2.evaluation.output_metrics import current_density, alpha
 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)
+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('--out_path', type=str, help='path to the target VTK files', required=False, default=None)
+    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 = nf2_path.replace('.nf2', '.vtk') if args.out_path is None else args.out_path
+
+    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'
+
+    output = SphericalOutput(nf2_path)
     result = output.load(radius_range, latitude_range, longitude_range, pixels_per_solRad, progress=True,
-                         metrics={'j': current_density, 'twist': twist})
+                         metrics={'j': current_density, 'alpha': alpha})
 
     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']}
+    metrics = result['metrics']
+    scalars = {'radius': radius, 'current_density': np.sum(metrics['j'] ** 2, -1) ** 0.5, 'alpha': metrics['alpha']}
     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]
+    save_vtk(out_path, coords, vectors, scalars)
diff --git a/nf2/evaluation/spherical/potential_vtk.py b/nf2/convert/potential_vtk.py
similarity index 80%
rename from nf2/evaluation/spherical/potential_vtk.py
rename to nf2/convert/potential_vtk.py
index 2a79e9b..406f10e 100644
--- a/nf2/evaluation/spherical/potential_vtk.py
+++ b/nf2/convert/potential_vtk.py
@@ -1,32 +1,27 @@
 import argparse
-import asyncio
-import glob
 import os
-import threading
 
 import numpy as np
 import pfsspy
 import torch
+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.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.data.util import cartesian_to_spherical, spherical_to_cartesian, vector_spherical_to_cartesian
 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('--Br', type=str, help='path to the source radial magnetic field file (full disk map)')
+parser.add_argument('--synoptic', type=str, help='path to the source radial magnetic field file (synoptic map)')
 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('--radius_range', nargs='+', type=float, default=(0.999, 1.5), 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)
+parser.add_argument('--pixels_per_solRad', type=int, default=64, required=False)
 
 args = parser.parse_args()
 out_path = args.out_path
@@ -44,11 +39,18 @@
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
+synoptic_map_br = Map(args.synoptic)
+synoptic_map_br = synoptic_map_br.resample([360 * 2, 180 * 2] * u.pix)
+
 mag_r_map = Map(args.Br)
-mag_r_map = mag_r_map.resample([360 * 4, 180 * 4] * u.pix)
-nrho = 9
+mag_r_map = mag_r_map.reproject_to(synoptic_map_br.wcs)
+
+nan_mask = ~np.isnan(mag_r_map.data)
+synoptic_map_br.data[nan_mask] = mag_r_map.data[nan_mask]
+
+nrho = 100
 rss = 2.5
-pfss_in = pfsspy.Input(mag_r_map, nrho, rss)
+pfss_in = pfsspy.Input(synoptic_map_br, nrho, rss)
 pfss_out = pfsspy.pfss(pfss_in)
 
 fig, axs = plt.subplots(10, 3, figsize=(10, 10))
@@ -68,11 +70,8 @@
 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),
@@ -84,9 +83,9 @@
 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
+x_min, x_max = -1.3, 1.3
+y_min, y_max = -1.3, 1.3
+z_min, z_max = -1.3, 1.3
 
 coords = np.stack(
     np.meshgrid(np.linspace(x_min, x_max, int((x_max - x_min) * pixels_per_solRad)),
@@ -102,7 +101,6 @@
 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,
@@ -112,5 +110,9 @@
 
 b = np.zeros(cube_shape + (3,))
 b[condition] = sub_b
+b = np.nan_to_num(b, nan=0)
+
+coords = np.stack(np.mgrid[0:b.shape[0], 0:b.shape[1], 0:b.shape[2]], -1).astype(np.int64)
+coords = (coords/ (b.shape[0] - 1) - 0.5) * 2 * 1.3
 
-save_vtk(vtk_path, {'B': b}, {'radius': radius})
+save_vtk(vtk_path, coords=coords, vectors={'B': b})
diff --git a/nf2/data/dataset.py b/nf2/data/dataset.py
index 6dc0d02..7482efd 100644
--- a/nf2/data/dataset.py
+++ b/nf2/data/dataset.py
@@ -176,9 +176,14 @@ 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)
+        if coord_range[2, 0] == 0:
+            z_range = np.linspace(coord_range[2, 0], coord_range[2, 1], n_slices, dtype=np.float32)
+        else:
+            z_range = np.linspace(0, coord_range[2, 1], n_slices - 1, dtype=np.float32)
+            z_range = np.concatenate([np.array([coord_range[2, 0]]), z_range])
         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),
+                                        z_range,
                                         indexing='ij'), -1)
         self.cube_shape = coords.shape[:-1]
         #
diff --git a/nf2/data/download_range.py b/nf2/data/download_range.py
index 74d205a..8099e6e 100644
--- a/nf2/data/download_range.py
+++ b/nf2/data/download_range.py
@@ -18,7 +18,7 @@ def main():
     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
+    sharp_num = args.sharp_num
     noaa_num = args.noaa_num
     t_start = args.t_start
     t_end = args.t_end
diff --git a/nf2/data/sst/mask.py b/nf2/data/sst/mask.py
new file mode 100644
index 0000000..f0fef20
--- /dev/null
+++ b/nf2/data/sst/mask.py
@@ -0,0 +1,32 @@
+import glob
+import os
+
+import numpy as np
+from astropy.io import fits
+from matplotlib import pyplot as plt
+
+fits_file = '/glade/work/rjarolim/data/SST/campaign_2023_1050/BLOS_panorama1050_6173_StkIQUV_rebin.fits'
+
+data_los = fits.getdata('/glade/work/rjarolim/data/SST/campaign_2023_1050/BLOS_panorama1050_6173_StkIQUV_rebin.fits')
+data_trv = fits.getdata('/glade/work/rjarolim/data/SST/campaign_2023_1050/BTRV_panorama1050_6173_StkIQUV_rebin.fits')
+header = fits.getheader(fits_file)
+
+data = (data_los - data_los[0, 0]) - (data_trv - data_trv[0, 0])
+mask = (data == 0).astype(np.float32)
+
+# save mask to a fits file
+hdu = fits.PrimaryHDU(data=mask, header=header)
+hdu.writeto('/glade/work/rjarolim/data/SST/campaign_2023_1050/mask.fits', overwrite=True)
+# save mask to a png file
+plt.imsave('/glade/work/rjarolim/data/SST/campaign_2023_1050/mask.png', mask, cmap='gray')
+plt.imsave('/glade/work/rjarolim/data/SST/campaign_2023_1050/data1.png', data, vmin=0, vmax=500)
+plt.imsave('/glade/work/rjarolim/data/SST/campaign_2023_1050/data2.png', fits.getdata('/glade/work/rjarolim/data/SST/campaign_2023_1050/BTRV_panorama1050_6173_StkIQUV_rebin.fits'), vmin=0, vmax=500)
+
+converted_path = '/glade/work/rjarolim/data/SST/campaign_2023_1050_converted'
+os.makedirs(converted_path, exist_ok=True)
+for f in glob.glob('/glade/work/rjarolim/data/SST/campaign_2023_1050/*.fits'):
+    d, h  = fits.getdata(f), fits.getheader(f)
+    d = np.flip(d, 0)
+    ft = os.path.join(converted_path, os.path.basename(f))
+    fits.writeto(ft, d, h, overwrite=True)
+
diff --git a/nf2/data/util.py b/nf2/data/util.py
index 5afd66b..61f2301 100644
--- a/nf2/data/util.py
+++ b/nf2/data/util.py
@@ -86,17 +86,20 @@ def cartesian_to_spherical(v, f=np):
 
 
 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_x, b_y, b_z = b[..., 0], b[..., 1], b[..., 2]
+
+    B_los = b_z
+    B_trv = (b_x ** 2 + b_y ** 2) ** 0.5
+    azi = f.arctan2(b_x, b_y)
+
     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_x = B_trv * f.sin(azi % f.pi) if ambiguous else B_trv * f.sin(azi)
+    B_y = B_trv * f.cos(azi % f.pi) 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
index 039874f..ccc3b5a 100644
--- a/nf2/evaluation/compute_metrics.py
+++ b/nf2/evaluation/compute_metrics.py
@@ -1,7 +1,8 @@
 import argparse
 
 from nf2.evaluation.metric import normalized_divergence, weighted_theta, theta_J, sigma_J
-from nf2.evaluation.output import current_density, CartesianOutput
+from nf2.evaluation.output import CartesianOutput
+from nf2.evaluation.output_metrics import current_density
 
 
 def main():
diff --git a/nf2/evaluation/muram/muram_to_vtk.py b/nf2/evaluation/muram/muram_to_vtk.py
deleted file mode 100644
index 05e9d91..0000000
--- a/nf2/evaluation/muram/muram_to_vtk.py
+++ /dev/null
@@ -1,33 +0,0 @@
-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/output.py b/nf2/evaluation/output.py
index 991b991..9813efc 100644
--- a/nf2/evaluation/output.py
+++ b/nf2/evaluation/output.py
@@ -10,102 +10,11 @@
 
 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
+from nf2.evaluation.metric import energy
+from nf2.evaluation.output_metrics import metric_mapping
+from nf2.train.model import VectorPotentialModel
+from nf2.train.transform import HeightTransformModel
+
 
 class BaseOutput:
 
@@ -117,7 +26,6 @@ def __init__(self, checkpoint, device=None):
         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
 
@@ -129,12 +37,13 @@ def G_per_dB(self):
     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}):
+    def load_coords(self, coords, batch_size=int(2 ** 12), progress=False, compute_jacobian=True, metrics=None):
         batch_size = batch_size * torch.cuda.device_count() if torch.cuda.is_available() else batch_size
+        metrics = metrics if metrics is not None else []
+
         def _load(coords):
             # normalize and to tensor
-            coords = torch.tensor(coords / self.spatial_norm, dtype=torch.float32)
+            coords = torch.tensor(coords, dtype=torch.float32)
             coords_shape = coords.shape
             coords = coords.reshape((-1, 3))
 
@@ -160,21 +69,26 @@ def _load(coords):
                 model_out['a'] = model_out['a'] * self.G_per_dB * self.m_per_ds
             return model_out
 
-        if compute_jacobian or self._requires_grad:
+        if self._requires_grad or compute_jacobian:
             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
+            if compute_jacobian:
+                jac_matrix = model_out['jac_matrix']
+                jac_matrix = jac_matrix * self.G_per_dB / self.m_per_ds
+                model_out['jac_matrix'] = jac_matrix
         else:
             with torch.no_grad():
-                return _load(coords)
+                model_out = _load(coords)
+
+        state = {**model_out, 'coords': coords}
+        metrics_out = {}
+        for key in metrics:
+            metric_out = metric_mapping[key](**state)
+            metrics_out.update(metric_out)
+            state.update(metric_out)
+
+        model_out['metrics'] = metrics_out
+        return model_out
 
 
 class CartesianOutput(BaseOutput):
@@ -191,7 +105,7 @@ def __init__(self, *args, **kwargs):
         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.time = None if self.wcs is None or len(self.wcs) >= 1 else parse(self.wcs[0].wcs.dateobs)
         self.data_config = self.state['data']
 
     def load_cube(self, height_range=None, Mm_per_pixel=None, **kwargs):
@@ -201,28 +115,27 @@ def load_cube(self, height_range=None, Mm_per_pixel=None, **kwargs):
                                                                                            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
+        ds_per_pixel = Mm_per_pixel / self.Mm_per_ds
 
         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)
-
+            np.meshgrid(np.linspace(x_min, x_max, np.round((x_max - x_min) / ds_per_pixel + 1).astype(int)),
+                        np.linspace(y_min, y_max, np.round((y_max - y_min) / ds_per_pixel + 1).astype(int)),
+                        np.linspace(z_min, z_max, np.round((z_max - z_min) / ds_per_pixel + 1).astype(int)),
+                        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):
+    def load_slice(self, z=0 * u.Mm, 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)
+            np.meshgrid(np.linspace(x_min, x_max, np.round((x_max - x_min) / self.ds_per_pixel + 1).astype(int)),
+                        np.linspace(y_min, y_max, np.round((y_max - y_min) / self.ds_per_pixel + 1).astype(int)),
+                        np.ones((1,), dtype=np.float32) * z.to_value(u.Mm) / self.Mm_per_pixel, indexing='ij'), -1)
         coords = coords[:, :, 0]
 
         model_out = self.load_coords(coords, **kwargs)
@@ -269,7 +182,7 @@ def trace(self, start_coords, direction=1, max_iterations=None, **kwargs):
         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
+                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)}
 
@@ -321,16 +234,23 @@ 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.transforms = self.state['transforms']
+        height_transforms = [t for t in self.transforms if isinstance(t, HeightTransformModel)]
+
+        assert len(height_transforms) == 1, 'Requires transform module!'
+        self.transform_module = height_transforms[0]
 
         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):
+        mapping_out = []
         for coord_range, height_mapping, ds_per_pixel in zip(self.coord_range_list, self.height_mapping_list,
                                                              self.ds_per_pixel_list):
+            if height_mapping is None:
+                continue
             x_min, x_max = coord_range[0]
             y_min, y_max = coord_range[1]
             z = height_mapping['z'] / self.Mm_per_ds
@@ -345,21 +265,23 @@ def load_height_mapping(self, **kwargs):
             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}
+            entry = {'height': z * self.Mm_per_ds, '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}
+            mapping_out.append(entry)
+        return mapping_out
 
     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 = torch.tensor(coords, 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 = []
+            cube = {}
             it = range(int(np.ceil(coords.shape[0] / batch_size)))
             it = tqdm(it) if progress else it
             for k in it:
@@ -370,14 +292,16 @@ def _load(coords, height_range):
                 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)
+                transformed_coords = self.transform_module({'coords': coord, 'height_range': height_range_batch})
 
-                cube += [transformed_coords.detach().cpu()]
+                for k, v in transformed_coords.items():
+                    if k not in cube:
+                        cube[k] = []
+                    cube[k] += [v.detach().cpu()]
 
-            cube = torch.cat(cube)
-            cube = cube.reshape(*coords_shape).numpy()
+            cube = {k: torch.cat(v).reshape(*coords_shape).numpy() for k, v in cube.items()}
 
-            return {'coords': cube}
+            return cube
 
         with torch.no_grad():
             return _load(coords, height_range)
@@ -390,7 +314,6 @@ def __init__(self, *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,
@@ -404,8 +327,8 @@ def load_spherical(self, radius_range: u.Quantity = None,
                 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)
+        scaled_coords = cartesian_coords * (1 * u.solRad / self.m_per_ds).to_value(u.dimensionless_unscaled)
+        model_out = self.load_coords(scaled_coords, **kwargs)
         return {**model_out, 'coords': cartesian_coords, 'spherical_coords': spherical_coords}
 
     def load(self,
@@ -452,15 +375,25 @@ def load(self,
                 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]
+        scaled_coords = coords * (1 * u.solRad / self.m_per_ds).to_value(u.dimensionless_unscaled)
+        sub_coords = scaled_coords[condition]
+
+        cube_shape = scaled_coords.shape[:-1]
         model_out = self.load_coords(sub_coords, **kwargs)
 
-        spherical_out = {'spherical_coords': spherical_coords, 'coords': coords}
+        spherical_out = {'spherical_coords': spherical_coords, 'coords': coords, 'metrics': {}}
+        metrics = model_out.pop('metrics')
+        for k, sub_v in metrics.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['metrics'][k] = volume
+
         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
+            if hasattr(sub_v, 'unit'):  # preserve units
                 volume = volume * sub_v.unit
             volume[condition] = sub_v
             spherical_out[k] = volume
@@ -471,7 +404,8 @@ def load(self,
 
     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)
+        scaled_coords = cartesian_coords * (1 * u.solRad / self.m_per_ds).to_value(u.dimensionless_unscaled)
+        model_out = self.load_coords(scaled_coords, **kwargs)
         model_out['spherical_coords'] = spherical_coords
         model_out['coords'] = cartesian_coords
         return model_out
diff --git a/nf2/evaluation/q_map.py b/nf2/evaluation/output_metrics.py
similarity index 56%
rename from nf2/evaluation/q_map.py
rename to nf2/evaluation/output_metrics.py
index 8464655..f7ef6fa 100644
--- a/nf2/evaluation/q_map.py
+++ b/nf2/evaluation/output_metrics.py
@@ -1,89 +1,168 @@
-import argparse
-import os.path
-
-import cupy
 import numpy as np
-from fastqslpy import FastQSL
-from fastqslpy.kernels import compileTraceBlineAdaptive
+from astropy import constants, units as u
+from astropy.nddata import block_reduce
 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)
+from nf2.data.util import cartesian_to_spherical
+from nf2.evaluation.energy import get_free_mag_energy
+from nf2.train.model import 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': 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_value(1 / u.Mm)
+    return {'b_nabla_bz': b_nabla_bz}
+
+
+def magnetic_helicity(b, a, **kwargs):
+    helicity = np.sum(a * b, axis=-1)
+    return {'magnetic_helicity': helicity}
+
+
+def alpha(b, jac_matrix, **kwargs):
+    j = calculate_current_from_jacobian(jac_matrix, f=np)
+    alpha = np.linalg.norm(j, axis=-1) / np.linalg.norm(b, axis=-1)
+    # threshold = np.linalg.norm(b, axis=-1).to_value(u.G) < 1  # set alpha to 0 for weak fields (coronal holes)
+    # alpha[threshold] = 0 * u.Mm ** -1
+    alpha = alpha.to(u.Mm ** -1)
+    return {'alpha': alpha}
+
+
+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': 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': 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)
+    b_los_trv_azi = np.stack([b_los, b_trv, azimuth], -1)
+    return {'b_los_trv_azi': b_los_trv_azi}
+
+
+def free_energy(b, **kwargs):
+    free_energy = get_free_mag_energy(b.to_value(u.G)) * u.erg * u.cm ** -3
+    return {'free_energy': free_energy}
+
+
+def squashing_factor(b, interp_ratio = 3, x_range=None, y_range=None, z_range=None, **kwargs):
+    # local imports for optional dependency
+    import cupy
+    from fastqslpy import FastQSL
+    from fastqslpy.kernels import compileTraceBlineAdaptive
 
     # 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)
-
+    Bx = b[..., 0].astype(np.float32)
+    By = b[..., 1].astype(np.float32)
+    Bz = 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)
+    # calc curl(B)
+    curBx_gpu = cupy.zeros_like(Bx_gpu)
+    curBy_gpu = cupy.zeros_like(By_gpu)
+    curBz_gpu = cupy.zeros_like(Bz_gpu)
+
+    curBx_gpu[:, 1:-1, 1:-1] = ((Bz_gpu[:, 2:, 1:-1] - Bz_gpu[:, 0:-2, 1:-1]) / 2.
+                                - (By_gpu[:, 1:-1, 2:] - By_gpu[:, 1:-1, 0:-2]) / 2)
+    curBy_gpu[1:-1, :, 1:-1] = ((Bx_gpu[1:-1, :, 2:] - Bx_gpu[1:-1, :, 0:-2]) / 2.
+                                - (Bz_gpu[2:, :, 1:-1] - Bz_gpu[0:-2, :, 1:-1]) / 2)
+    curBz_gpu[1:-1, 1:-1, :] = ((By_gpu[2:, 1:-1, :] - By_gpu[0:-2, 1:-1, :]) / 2.
+                                - (Bx_gpu[1:-1, 2:, :] - Bx_gpu[1:-1, 0:-2, :]) / 2)
+
+    # take care of z=0
+    curBx_gpu[1:-1, 1:-1, 0] = ((Bz_gpu[1:-1, 2:, 0] - Bz_gpu[1:-1, 0:-2, 0]) / 2.
+                                - (-3. * By_gpu[1:-1, 1:-1, 0] + 4. * By_gpu[1:-1, 1:-1, 1] - By_gpu[1:-1, 1:-1,
+                                                                                              2]) / 2)
+    curBy_gpu[1:-1, 1:-1, 0] = ((-3. * Bx_gpu[1:-1, 1:-1, 0] + 4. * Bx_gpu[1:-1, 1:-1, 1] - Bx_gpu[1:-1, 1:-1, 2]) / 2.
+                                - (Bz_gpu[2:, 1:-1, 0] - Bz_gpu[0:-2, 1:-1, 0]) / 2)
 
     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]]
+    x_range = [0, Bx.shape[0]] if x_range is None else x_range
+    y_range = [0, Bx.shape[1]] if y_range is None else y_range
+    z_range = [0, Bx.shape[2]] if z_range is None else z_range
     # 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()
 
@@ -91,19 +170,14 @@
     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'):
 
@@ -121,7 +195,7 @@
         # 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,
+                          curBx_gpu, curBy_gpu, curBz_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,
@@ -234,16 +308,21 @@
 
         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"]})
+    Qube = block_reduce(Qube, (interp_ratio, interp_ratio, interp_ratio), func=np.mean)
+    Twube = block_reduce(Twube, (interp_ratio, interp_ratio, interp_ratio), func=np.mean)
+
+    return {"q": Qube, "twist": Twube}
+
+
+metric_mapping = {
+    'j': current_density,
+    'b_nabla_bz': b_nabla_bz,
+    'alpha': alpha,
+    'spherical_energy_gradient': spherical_energy_gradient,
+    'energy_gradient': energy_gradient,
+    'magnetic_helicity': magnetic_helicity,
+    'los_trv_azi': los_trv_azi,
+    'free_energy': free_energy,
+    'squashing_factor': squashing_factor
+}
diff --git a/nf2/evaluation/sharp/convert_series.py b/nf2/evaluation/sharp/convert_series.py
index c0d39d3..3cdb801 100644
--- a/nf2/evaluation/sharp/convert_series.py
+++ b/nf2/evaluation/sharp/convert_series.py
@@ -2,84 +2,63 @@
 import glob
 import os
 import pickle
-from copy import copy
 
 import numpy as np
+import torch.cuda
 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.metric import energy, vector_norm, divergence, theta_J, sigma_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)
+    # list of nf2 paths
+    parser.add_argument('--nf2_dir', type=str, help='path to the directory with the NF2 files', required=True, nargs='+')
     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('--height_range', type=float, help='height range', required=False, nargs=2, default=None)
     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)
+    parser.add_argument('--batch_size', type=int, help='batch size', required=False, default=int(2 ** 13))
     args = parser.parse_args()
 
-    nf2_files = sorted(glob.glob(args.nf2_path))
+    nf2_files = [sorted(glob.glob(f)) for f in args.nf2_dir]
+    nf2_files = [f for files in nf2_files for f in files] # flatten list
     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)
+    batch_size = args.batch_size * torch.cuda.device_count() if torch.cuda.is_available() else args.batch_size
+    convert_nf2_series(nf2_files, result_path,
+                       height_range=args.height_range,
+                       Mm_per_pixel=args.Mm_per_pixel, batch_size=batch_size)
 
 
 def evaluate_nf2(nf2_file, **kwargs):
     out = CartesianOutput(nf2_file)
-    res = out.load_cube(**kwargs)
+    res = out.load_cube(metrics=['j'], **kwargs)
 
     b = res['b']
-    j = res['j']
+    j = res['metrics']['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,
+            'free_energy': free_me.sum() * Mm_per_pixel ** 3
         },
         'metrics': {
             'divergence': (np.abs(divergence(b)) / vector_norm(b)).mean(),
             'jxb': vector_norm(jxb).mean(),
-            'theta': theta
+            'theta': theta,
         },
         'maps': {
             'b_0': b[:, :, 0, 2],  # bottom boundary
@@ -101,106 +80,31 @@ def evaluate_nf2(nf2_file, **kwargs):
     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 = []
+def convert_nf2_series(nf2_paths, result_path, overwrite=True, **kwargs):
     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)
+        save_file = os.path.join(result_path, os.path.basename(nf2_file).replace('.nf2', '.pkl'))
+        if os.path.exists(save_file) and overwrite is False:
             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': []}
+    integrated_quantities = {'energy': [], 'free_energy': []}
     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:
+    for f in tqdm(series_results, desc='Loading files'):
         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'])
diff --git a/nf2/evaluation/sharp/ensemble.py b/nf2/evaluation/sharp/ensemble.py
index 4e2da59..58e2ea2 100644
--- a/nf2/evaluation/sharp/ensemble.py
+++ b/nf2/evaluation/sharp/ensemble.py
@@ -14,27 +14,24 @@
 
 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('--ref_path', type=str, help='path to the reference map')
 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()
 
+result_path = args.result_path
+os.makedirs(result_path, exist_ok=True)
+
 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)
+    out = model.load_cube(Mm_per_pixel=args.Mm_per_pixel, metrics=['j'])
     b_ensemble.append(out['b'])
-    j_ensemble.append(out['j'])
+    j_ensemble.append(out['metrics']['j'])
 
 Mm_per_pixel = out['Mm_per_pixel']
 
@@ -50,60 +47,36 @@
 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))
+fig, axs = plt.subplots(1, 3, figsize=(12, 2))
 
+ax = axs[0]
 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.colorbar(im, cax=cax, orientation='vertical', label='B$_z$ [G]')
 
-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)
+ax = axs[1]
+im = ax.imshow(b_uncertainty.mean(2).T, cmap='viridis', origin='lower', extent=extent_xy, vmin=0, vmax=20)
 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]')
+fig.colorbar(im, cax=cax, orientation='vertical', label='$\delta$B [G]')
 
-ax = axs[1]
-im = ax.imshow(b_uncertainty.mean(1).T, cmap='viridis', origin='lower', extent=extent_xz)
+ax = axs[2]
+b_norm = np.linalg.norm(mean_b_vector[0], axis=-1)
+relative_uncertainty = b_uncertainty / (b_norm + 1e-6) * 100
+im = ax.imshow(np.nanmean(relative_uncertainty, 2).T,
+               cmap='cividis', origin='lower', extent=extent_xy, vmin=0, vmax=30)
 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]')
-
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\delta$B$_{\text{relative}}$ [%]')
 
 plt.tight_layout()
-plt.savefig('/glade/work/rjarolim/nf2/sharp/13664/ensemble/overview.png', dpi=300, transparent=True)
+plt.savefig(os.path.join(result_path, 'uncertainty.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, :]]
+fig, axs = plt.subplots(1, 6, figsize=(21, 3))
 
 axs = np.ravel(axs)
 for i, (f, ax) in enumerate(zip(runs, axs[:-1])):
@@ -124,5 +97,6 @@
 ax.set_title('Ensemble Std')
 
 plt.tight_layout()
-plt.savefig(os.path.join(args.result_path, 'j_overview.png'), dpi=300, transparent=True)
+
+plt.savefig(os.path.join(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
index c71d66f..2a74867 100644
--- a/nf2/evaluation/sharp/flares.py
+++ b/nf2/evaluation/sharp/flares.py
@@ -9,8 +9,7 @@
 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 matplotlib.colors import LogNorm, Normalize
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from sunpy.map import Map
 
@@ -21,11 +20,14 @@
 
 
 class _F:
+
     def __init__(self, ref_wcs):
         self.ref_wcs = ref_wcs
 
     def func(self, file):
         s_map = Map(file)
+        if s_map.meta['QUALITY'] != 0:
+            return None
         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)
@@ -34,158 +36,136 @@ def func(self, file):
 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)
+    integrated_euv = np.array([m.data for m in reprojected_maps if m is not None]).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)
-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()
+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=peak_time,
+                             dir=result_path, client=client, channel=94)
+
+    filter_dates = (dates > (start_time - timedelta(minutes=12))) & \
+                   (dates < (peak_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(Mm_per_pixel=0.72, progress=True)
+    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(Mm_per_pixel=0.72, progress=True)
+    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
+    released_energy = -np.clip(energy_2 - energy_1, 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 < 1e10] = np.nan
+    total_released_energy = (released_energy.sum() * Mm_per_pix ** 3).to_value(u.erg)
+
+    # gained energy
+    gained_energy = np.clip(energy_2 - energy_1, a_min=0, a_max=None)
+    gained_energy_map = (gained_energy.sum(2) * Mm_per_pix).to_value(u.erg / u.cm ** 2)
+    gained_energy_map[gained_energy_map < 1e10] = np.nan
+    total_gained_energy = (gained_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)
+    energy_diff_map = (energy_2 - energy_1).sum(2) * Mm_per_pix
+    energy_diff_map = energy_diff_map.to_value(u.erg / u.cm ** 2)
+
+    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 * 1e-31:.02f} 10^31 erg\n')
+        f.write(f'Total gained energy: {total_gained_energy * 1e-31:.02f} 10^31 erg\n')
+        f.write(f'Total energy difference: {energy_diff * 1e-31:.02f} 10^31 erg\n')
+        f.write(f'Total free energy difference: {free_energy_diff * 1e-31:.02f} 10^31 erg\n')
+
+    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_1.sum(2) * Mm_per_pix).to_value(u.erg / u.cm ** 2)
+    free_energy_map[free_energy_map < 1e11] = 1e11
+    im = ax.imshow(free_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=LogNorm(1e11, 5e13))
+    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(energy_diff_map.T, origin='lower', cmap='seismic', extent=extent,
+                   norm=Normalize(vmin=-1e13, vmax=1e13))
+    # norm=SymLogNorm(1e12, vmin=-1e14, vmax=1e14))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="3%", pad=0.05)
+    fig.colorbar(im, cax=cax, label=r'$\Delta \text{E}$ [erg / cm$^2$]')
+
+    [ax.set_xlim(100, 400) for ax in axs]
+
+    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()
+
+    fig, axs = plt.subplots(1, 2, figsize=(7, 3))
+
+    ax = axs[0]
+    im = ax.imshow(gained_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=LogNorm(1e10, 1e13))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="3%", pad=0.05)
+    fig.colorbar(im, cax=cax, label=r'$\Delta \text{E}_\text{+}$ [erg / cm$^2$]')
+    ax.set_title('Gained Energy')
+
+    ax = axs[1]
+    im = ax.imshow(released_energy_map.T, origin='lower', cmap='jet', extent=extent, norm=LogNorm(1e10, 1e13))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="3%", pad=0.05)
+    fig.colorbar(im, cax=cax, label=r'$\Delta \text{E}_\text{-}$ [erg / cm$^2$]')
+    ax.set_title('Released Energy')
+
+    [ax.set_xlim(100, 400) for ax in axs]
+
+    plt.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_current_series.py b/nf2/evaluation/sharp/plot_current_series.py
new file mode 100644
index 0000000..4b0f337
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_current_series.py
@@ -0,0 +1,134 @@
+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
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.sharp.convert_series import load_results
+
+# 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)
+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
+
+times = output['times']
+wcs = output['wcs']
+maps = output['maps']
+Mm_per_pixel = output['Mm_per_pixel']
+
+j_norm = LogNorm()
+
+selected_times = [
+    datetime(2024, 5, 10, 5, 0, 0),
+    datetime(2024, 5, 10, 6, 0, 0),
+    datetime(2024, 5, 10, 7, 0, 0),
+    datetime(2024, 5, 10, 8, 0, 0),
+    datetime(2024, 5, 10, 9, 0, 0),
+    datetime(2024, 5, 10, 10, 0, 0),
+]
+
+cm = copy(plt.get_cmap('sdoaia131'))
+cm.set_bad('black')
+
+fig, axs = plt.subplots(len(selected_times) + 1, 2, figsize=(5.6, 2 * (len(selected_times) + 1)), sharex=True)
+
+for i, time in tqdm(enumerate(selected_times), total=len(selected_times)):
+    nf2_cond = np.argmin(np.abs(times - time))
+    current_density_map = maps['current_density_map'][nf2_cond]
+    current_density_map = current_density_map.to_value(u.G * u.cm * u.s ** -1)
+    map_wcs = wcs[nf2_cond][0]
+    #
+    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(map_wcs)
+    #
+    extent = np.array([0, current_density_map.shape[0],
+                       0, current_density_map.shape[1]]) * Mm_per_pixel
+    #
+    row = axs[i, :]
+    #
+    ax = row[0]
+    cd_im = ax.imshow(current_density_map.T, origin='lower',
+                      cmap='inferno', extent=extent, norm=j_norm)
+    #
+    ax = row[1]
+    euv_im = ax.imshow(euv_map.data / exposure, origin='lower',
+                       cmap=cm, extent=extent, norm=LogNorm(vmin=5, vmax=1e3))
+    # add time as label on the right in bold
+    ax.set_ylabel(time.isoformat(' ', timespec='minutes'), fontweight='bold')
+    ax.yaxis.set_label_position("right")
+
+# plot difference map
+first_cond = np.argmin(np.abs(times - selected_times[0]))
+current_density_map = maps['current_density_map'][first_cond]
+first_density_map = current_density_map.to_value(u.G * u.cm * u.s ** -1)
+
+last_cond = np.argmin(np.abs(times - selected_times[-1]))
+current_density_map = maps['current_density_map'][last_cond]
+last_density_map = current_density_map.to_value(u.G * u.cm * u.s ** -1)
+
+diff_density_map = last_density_map - first_density_map
+
+extent = np.array([0, current_density_map.shape[0],
+                   0, current_density_map.shape[1]]) * Mm_per_pixel
+ax = axs[-1, 0]
+im = ax.imshow(diff_density_map.T, origin='lower', extent=extent, vmin=-5e12, vmax=5e12, cmap='RdBu_r')
+
+axs[-1, 1].set_axis_off()
+axs[-2, 1].set_xlabel('X [Mm]')
+
+divider = make_axes_locatable(axs[0, 0])
+cax = divider.append_axes("top", size="5%", pad=0.05)
+fig.colorbar(cd_im, cax=cax, label='Current Density [G cm / s]', orientation='horizontal')
+cax.xaxis.set_ticks_position('top')
+cax.xaxis.set_label_position('top')
+
+divider = make_axes_locatable(axs[0, 1])
+cax = divider.append_axes("top", size="5%", pad=0.05)
+fig.colorbar(euv_im, cax=cax, label='SDO/AIA 131 $\AA$ [DN / s]', orientation='horizontal')
+cax.xaxis.set_ticks_position('top')
+cax.xaxis.set_label_position('top')
+
+[ax.set_xlim(100, 300) for ax in np.ravel(axs)]
+[ax.set_ylim(20, 150) for ax in np.ravel(axs)]
+
+[plt.setp(ax.get_xticklabels(), visible=False) for ax in np.ravel(axs[:-1, :])]
+[plt.setp(ax.get_yticklabels(), visible=False) for ax in np.ravel(axs[:, 1:])]
+
+[ax.set_xlabel('X [Mm]') for ax in np.ravel(axs[-1, :])]
+[ax.set_ylabel('Y [Mm]') for ax in np.ravel(axs[:, 0])]
+
+fig.tight_layout(pad=0.1)
+fig.savefig(result_path, dpi=300)
+plt.close()
+
+fig, ax = plt.subplots(1, 1, figsize=(5.6 / 2, 2))
+
+ax.set_axis_off()
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="5%", pad=0.05)
+fig.colorbar(im, cax=cax, label='$\Delta$ Current Density\n[G cm / s]', orientation='vertical')
+
+fig.tight_layout()
+fig.savefig(result_path.replace('.jpg', '_colorbar.png'), dpi=300, transparent=True)
+plt.close()
\ No newline at end of file
diff --git a/nf2/evaluation/sharp/plot_euv_lightcurve.py b/nf2/evaluation/sharp/plot_euv_lightcurve.py
new file mode 100644
index 0000000..ec64b09
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_euv_lightcurve.py
@@ -0,0 +1,68 @@
+import argparse
+import glob
+import os
+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
+from matplotlib.dates import DateFormatter
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.sharp.convert_series import load_results
+
+
+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)
+    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']
+
+    euv_lightcurve = []
+    cond = (times > datetime(2024, 5, 9, 15)) & (times < datetime(2024, 5, 10, 12))
+    times = times[cond]
+    wcs = wcs[cond]
+
+    for i, time in tqdm(enumerate(times), total=len(times)):
+        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])
+        euv_lightcurve.append(euv_map.data.mean() / exposure)
+
+    date_format = DateFormatter('%d-%H:%M')
+
+    fig, ax = plt.subplots(1, 1, figsize=(5, 3))
+
+    ax.xaxis_date()
+    ax.xaxis.set_major_formatter(date_format)
+
+    fig.autofmt_xdate()
+
+    ax.plot(times, euv_lightcurve)
+    ax.set_ylabel('SDO/AIA 131 $\AA$')
+    ax.set_yscale('log')
+
+    plt.tight_layout()
+    plt.savefig(os.path.join(result_path, 'euv_lightcurve.jpg'), dpi=300)
+    plt.close()
+
+
+# if __name__ == '__main__':
+#     main()
diff --git a/nf2/evaluation/sharp/plot_integrated_energy.py b/nf2/evaluation/sharp/plot_integrated_energy.py
new file mode 100644
index 0000000..05fbfc3
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_integrated_energy.py
@@ -0,0 +1,242 @@
+import argparse
+import glob
+import os
+from datetime import datetime, timedelta
+
+import numpy as np
+import pandas as pd
+from astropy import units as u
+from matplotlib import pyplot as plt
+from matplotlib.dates import DateFormatter
+from sunpy.map import Map
+from sunpy.net import Fido
+from sunpy.timeseries import TimeSeries
+from tqdm import tqdm
+
+from nf2.evaluation.sharp.convert_series import load_results
+from sunpy.net import attrs as a
+
+# 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)
+parser.add_argument('--sharp_path', type=str, help='path to the output directory', required=False, default=None)
+args = parser.parse_args()
+
+Br = sorted(glob.glob(os.path.join(args.sharp_path, '*Br.fits')))
+Bt = sorted(glob.glob(os.path.join(args.sharp_path, '*Bt.fits')))
+Bp = sorted(glob.glob(os.path.join(args.sharp_path, '*Bp.fits')))
+
+o_unsigned_flux = []
+for Br_file, Bt_file, Bp_file in tqdm(zip(Br, Bt, Bp), total=len(Br)):
+    # vector_B = np.stack([Map(Br_file).data, Map(Bt_file).data, Map(Bp_file).data], axis=-1)
+    # flux = np.linalg.norm(vector_B, axis=-1).sum()
+    flux = np.abs(Map(Br_file).data).sum() * (0.36 * u.Mm).to_value(u.cm) ** 2
+    o_unsigned_flux.append(flux)
+
+o_unsigned_flux = np.array(o_unsigned_flux)
+print(f'MIN/MAX unsigned flux: {o_unsigned_flux.min():.2e}/{o_unsigned_flux.max():.2e}')
+output = load_results(args.pkl_path)
+
+o_times = output['times']
+integrated_quantities = output['integrated_quantities']
+result_path = args.result_path
+Mm_per_pixel = output['Mm_per_pixel']
+
+os.makedirs(result_path, exist_ok=True)
+
+o_free_energy = integrated_quantities['free_energy'].to_value(u.erg) * 1e-32
+o_energy = integrated_quantities['energy'].to_value(u.erg) * 1e-32
+o_current = output['maps']['current_density_map'].sum((1, 2)) * Mm_per_pixel ** 2
+
+date_format = DateFormatter('%d-%H:%M')
+
+# fill gaps with NaNs
+df = pd.DataFrame({'time': o_times, 'unsigned_flux': o_unsigned_flux,
+                   'free_energy': o_free_energy, 'energy': o_energy, 'current': o_current})
+df = df.set_index('time')
+df = df.resample('12min').mean()
+df = df.rolling('1H').mean()
+
+# load GOES data
+result_goes = Fido.search(a.Time(o_times[0], o_times[-1]),
+                          a.Instrument.xrs & a.goes.SatelliteNumber(16) & a.Resolution("avg1m"))
+file_goes = Fido.fetch(result_goes)
+goes = TimeSeries(file_goes, source='XRS', concatenate=True)
+
+dt = 12 * 60  # 12 minutes in seconds
+
+times = df.index
+unsigned_flux = df['unsigned_flux'].values
+free_energy = df['free_energy'].values
+energy = df['energy'].values
+current = df['current'].values
+
+fig, axs = plt.subplots(5, 1, figsize=(12, 8))
+
+# make date axis
+for ax in axs:
+    ax.xaxis_date()
+    ax.set_xlim(datetime(2024, 5, 7), times[-1])
+    ax.xaxis.set_major_formatter(date_format)
+
+fig.autofmt_xdate()
+
+ax = axs[0]
+color = 'tab:blue'
+ax.plot(times, unsigned_flux * 1e-23, color=color)
+ax.set_ylabel('$|B_z|$\n[$10^{23}$ Mx]', color=color)
+ax.spines['left'].set_color(color)
+ax.yaxis.label.set_color(color)
+ax.tick_params(axis='y', colors=color)
+
+twin_ax = ax.twinx()
+color = 'tab:red'
+twin_ax.plot(times, np.gradient(unsigned_flux) / dt * 1e-18, color=color)
+twin_ax.set_ylabel('$\mathrm{d} |B_z| / \mathrm{dt}$\n[$10^{18}$ Mx/s]', color=color)
+twin_ax.spines['right'].set_color(color)
+twin_ax.yaxis.label.set_color(color)
+twin_ax.tick_params(axis='y', colors=color)
+twin_ax.axhline(0, color='black', linestyle='dotted')
+
+ax = axs[1]
+color = 'tab:blue'
+ax.plot(times, energy)
+ax.set_ylabel('$E$\n[$10^{32}$ erg]', color=color)
+ax.spines['left'].set_color(color)
+ax.yaxis.label.set_color(color)
+ax.tick_params(axis='y', colors=color)
+
+twin_ax = ax.twinx()
+color = 'tab:red'
+twin_ax.plot(times, np.gradient(energy) / dt * 1e3, color=color)
+twin_ax.set_ylabel('$\mathrm{d}E / \mathrm{dt}$\n [$10^{29}$ erg/s]', color=color)
+twin_ax.spines['right'].set_color(color)
+twin_ax.yaxis.label.set_color(color)
+twin_ax.tick_params(axis='y', colors=color)
+twin_ax.axhline(0, color='black', linestyle='dotted')
+
+ax = axs[2]
+color = 'tab:blue'
+ax.plot(times, free_energy, color=color)
+ax.set_ylabel(r'$E_\text{free}$'+'\n[$10^{32}$ erg]', color=color)
+ax.spines['left'].set_color(color)
+ax.yaxis.label.set_color(color)
+ax.tick_params(axis='y', colors=color)
+
+twin_ax = ax.twinx()
+color = 'tab:red'
+twin_ax.plot(times, np.gradient(free_energy) / dt * 1e3, color=color)
+twin_ax.set_ylabel(r'$\mathrm{d} E_\text{free} / \mathrm{dt}$' + '\n [$10^{29}$ erg/s]', color=color)
+twin_ax.spines['right'].set_color(color)
+twin_ax.yaxis.label.set_color(color)
+twin_ax.tick_params(axis='y', colors=color)
+twin_ax.axhline(0, color='black', linestyle='dotted')
+
+energy_ratio = free_energy / energy
+ax = axs[3]
+color = 'tab:blue'
+ax.plot(times, energy_ratio, color=color)
+ax.set_ylabel(r'$E_\text{free} / E$', color=color)
+ax.spines['left'].set_color(color)
+ax.yaxis.label.set_color(color)
+ax.tick_params(axis='y', colors=color)
+
+twin_ax = ax.twinx()
+color = 'tab:red'
+twin_ax.plot(times, np.gradient(energy_ratio) / dt * 1e6, color=color)
+twin_ax.set_ylabel(r'$\mathrm{d} (E_\text{free} / E) / \mathrm{dt}$' + '\n [$10^{-6}$ / s]', color=color)
+twin_ax.spines['right'].set_color(color)
+twin_ax.yaxis.label.set_color(color)
+twin_ax.tick_params(axis='y', colors=color)
+twin_ax.axhline(-.5, color='gray', linestyle='--')
+twin_ax.axhline(0, color='black', linestyle='dotted')
+
+# plot GOES
+ax = axs[4]
+ax.plot(goes.data.index, goes.data['xrsb'], color='black')
+ax.set_yscale("log")
+ax.set_ylim(1e-6, 1e-3)
+ax.set_ylabel('GOES 1-8 $\AA$\n' + r'[W m$^{-2}$]')
+
+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')
+
+
+m_flare_times = [
+    (datetime(2024, 5, 7, 8, 18), datetime(2024, 5, 7, 8, 40)),  # M1.3
+    (datetime(2024, 5, 7, 11, 40), datetime(2024, 5, 7, 12, 7)),  # M2.5
+    (datetime(2024, 5, 7, 20, 18), datetime(2024, 5, 7, 20, 34)),  # M2.1
+    (datetime(2024, 5, 7, 21, 42), datetime(2024, 5, 7, 22, 21)),  # M3.3
+    #
+    (datetime(2024, 5, 8, 2, 16), datetime(2024, 5, 8, 2, 36)),  # M3.4
+    (datetime(2024, 5, 8, 3, 19), datetime(2024, 5, 8, 3, 38)),  # M1.9
+    (datetime(2024, 5, 8, 4, 37), datetime(2024, 5, 8, 5, 32)),  # M3.6
+    (datetime(2024, 5, 8, 6, 44), datetime(2024, 5, 8, 7, 10)),  # M7.2
+    (datetime(2024, 5, 8, 11, 26), datetime(2024, 5, 8, 12, 22)),  # M8.7
+    # (datetime(2024, 5, 8, 17, 32), datetime(2024, 5, 8, 18, 0)),  # M7.9
+    # (datetime(2024, 5, 8, 19, 15), datetime(2024, 5, 8, 19, 29)),  # M2.0
+    #
+    (datetime(2024, 5, 9, 3, 7), datetime(2024, 5, 9, 3, 23)),  # M4.0
+    (datetime(2024, 5, 9, 3, 23), datetime(2024, 5, 9, 3, 49)),  # M4.6
+    (datetime(2024, 5, 9, 4, 44), datetime(2024, 5, 9, 4, 55)),  # M1.7
+    (datetime(2024, 5, 9, 6, 3), datetime(2024, 5, 9, 6, 31)),  # M2.3
+    (datetime(2024, 5, 9, 6, 24), datetime(2024, 5, 9, 6, 31)),  # M2.5
+    (datetime(2024, 5, 9, 11, 52), datetime(2024, 5, 9, 12, 2)),  # M3.1
+    (datetime(2024, 5, 9, 12, 5), datetime(2024, 5, 9, 12, 20)),  # M2.9
+    (datetime(2024, 5, 9, 13, 16), datetime(2024, 5, 9, 13, 29)),  # M3.7
+    # (datetime(2024, 5, 8, 22, 4), datetime(2024, 5, 8, 23, 30)),  # M9.9
+    # (datetime(2024, 5, 8, 22, 5), datetime(2024, 5, 8, 22, 45)),  # M9.7
+    (datetime(2024, 5, 9, 22, 24), datetime(2024, 5, 9, 22, 47)),  # M2.6
+    (datetime(2024, 5, 9, 23, 44), datetime(2024, 5, 9, 23, 55)),  # M1.6
+    #
+    (datetime(2024, 5, 10, 0, 10), datetime(2024, 5, 10, 0, 22)),  # M1.3
+    (datetime(2024, 5, 10, 3, 15), datetime(2024, 5, 10, 3, 40)),  # M1.4
+    (datetime(2024, 5, 10, 10, 10), datetime(2024, 5, 10, 10, 19)),  # M2.2
+    (datetime(2024, 5, 10, 13, 58), datetime(2024, 5, 10, 14, 28)),  # M6.0
+    (datetime(2024, 5, 10, 18, 26), datetime(2024, 5, 10, 18, 38)),  # M1.1
+    (datetime(2024, 5, 10, 18, 38), datetime(2024, 5, 10, 19, 7)),  # M1.8
+    (datetime(2024, 5, 10, 18, 57), datetime(2024, 5, 10, 19, 10)),  # M2.0
+    (datetime(2024, 5, 10, 19, 35), datetime(2024, 5, 10, 19, 56)),  # M1.1
+    (datetime(2024, 5, 10, 19, 56), datetime(2024, 5, 10, 20, 19)),  # M1.9
+    (datetime(2024, 5, 10, 20, 59), datetime(2024, 5, 10, 21, 12)),  # M3.8
+    #
+    (datetime(2024, 5, 11, 4, 28), datetime(2024, 5, 11, 4, 51)),  # M1.4
+]
+
+x_flares = [
+    # (datetime(2024, 5, 8, 21, 8), datetime(2024, 5, 8, 21, 58)),  # X1.0
+    (datetime(2024, 5, 8, 4, 37), datetime(2024, 5, 8, 5, 32)),  # X1.0
+    (datetime(2024, 5, 9, 8, 45), datetime(2024, 5, 9, 9, 36)),  # X2.2
+    (datetime(2024, 5, 9, 17, 23), datetime(2024, 5, 9, 18, 0)),  # X1.1
+    (datetime(2024, 5, 10, 6, 27), datetime(2024, 5, 10, 7, 6)),  # X4.0
+    (datetime(2024, 5, 11, 1, 10), datetime(2024, 5, 11, 1, 39)),  # X5.8
+]
+
+for start, end in m_flare_times:
+    [ax.axvspan(start, end, color='blue', alpha=0.2) for ax in axs]
+
+for start, end in x_flares:
+    [ax.axvspan(start, end, color='orange', alpha=1) for ax in axs]
+
+nans = np.isnan(unsigned_flux)
+# pairs of ranges
+ranges = np.where(np.diff(nans) != 0)[0].reshape(-1, 2)
+for start, end in ranges:
+    [ax.axvspan(times[start], times[end], color='gray', alpha=0.3) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(result_path, 'integrated_energy.jpg'), dpi=300)
+plt.close()
+
+energy_change = np.gradient(energy_ratio) / dt
+for x_flare in x_flares:
+    start, end = x_flare
+    cond = (times > start) & (times < end)
+    e = energy_change[cond]
+    print(f"MEAN ({start.isoformat(' ', timespec='minutes')}): {e.mean() * 1e6:.2f}")
+
+print(f'MEAN: {np.nanmean(energy_change) * 1e6:.2f}')
\ No newline at end of file
diff --git a/nf2/evaluation/sharp/plot_integrated_quantities.py b/nf2/evaluation/sharp/plot_integrated_quantities.py
index bdcb7ae..899fa94 100644
--- a/nf2/evaluation/sharp/plot_integrated_quantities.py
+++ b/nf2/evaluation/sharp/plot_integrated_quantities.py
@@ -75,7 +75,6 @@ def _plot_integrated_qunatities(times, integrated_quantities, height_distributio
 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
@@ -84,7 +83,7 @@ def _plot_integrated_qunatities(times, integrated_quantities, height_distributio
 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))
+                cmap=cm.get_cmap('jet'), )
 ax.set_ylabel('Height\n[Mm]')
 ax.set_ylim(0, 20)
 # add colorbar
@@ -94,13 +93,11 @@ def _plot_integrated_qunatities(times, integrated_quantities, height_distributio
 # 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.plot(goes.data.index, goes.data['xrsb'], label='1-8 $\AA$')
 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')
@@ -116,7 +113,8 @@ def _plot_integrated_qunatities(times, integrated_quantities, height_distributio
     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')
-
+    # plot 60 degree line
+    ax.axvline(x=datetime(2024, 5, 10, 15, 34), linestyle='solid', c='magenta')
 
 # labels = ['A', 'B', 'C', 'M', 'X']
 # centers = np.logspace(-7.5, -3.5, len(labels))
diff --git a/nf2/evaluation/sharp/plot_maps.py b/nf2/evaluation/sharp/plot_maps.py
index abb9f7b..c892480 100644
--- a/nf2/evaluation/sharp/plot_maps.py
+++ b/nf2/evaluation/sharp/plot_maps.py
@@ -53,6 +53,7 @@ def main():
         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)
+        free_energy_map[free_energy_map < 1e9] = 1e9
 
         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
@@ -95,7 +96,7 @@ def main():
 
     fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
     ax.set_axis_off()
-    cbar = fig.colorbar(b0_im, ax=ax, label='[G]')
+    cbar = fig.colorbar(b0_im, ax=ax, label='$B_z$ [G]')
     # set ticks
     cbar.set_ticks([-2000, -1000, 0, 1000, 2000])
     plt.tight_layout(pad=0)
@@ -104,14 +105,14 @@ def main():
 
     fig, ax = plt.subplots(1, 1, figsize=(3, 1.5))
     ax.set_axis_off()
-    fig.colorbar(euv_im, ax=ax, label='[DN / s]')
+    fig.colorbar(euv_im, ax=ax, label='AIA 131 $\AA$ [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]')
+    fig.colorbar(cd_im, ax=ax, label='$J$ [G cm / s]')
     plt.tight_layout(pad=0)
     plt.savefig(os.path.join(result_path, f'j_colorbar.jpg'), dpi=300)
     plt.close()
diff --git a/nf2/evaluation/sharp/plot_quality.py b/nf2/evaluation/sharp/plot_quality.py
new file mode 100644
index 0000000..ebe252c
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_quality.py
@@ -0,0 +1,160 @@
+import argparse
+import glob
+import multiprocessing
+import os
+from datetime import datetime
+from multiprocessing import Pool
+
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib.dates import DateFormatter
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.output import CartesianOutput
+from nf2.evaluation.sharp.convert_series import load_results
+
+from astropy import units as u
+
+
+
+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('--fits_path', type=str, help='path to the directory with the fits files.', required=False, default=None)
+parser.add_argument('--nf2_dir', type=str, help='path to the directory with the NF2 files', required=True, nargs='+')
+parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+args = parser.parse_args()
+
+result_path = args.result_path
+
+b_r_files = sorted(glob.glob(os.path.join(args.fits_path, '*Br.fits')))
+b_t_files = sorted(glob.glob(os.path.join(args.fits_path, '*Bt.fits')))
+b_p_files = sorted(glob.glob(os.path.join(args.fits_path, '*Bp.fits')))
+b_r_err_files = sorted(glob.glob(os.path.join(args.fits_path, '*Br_err.fits')))
+b_t_err_files = sorted(glob.glob(os.path.join(args.fits_path, '*Bt_err.fits')))
+b_p_err_files = sorted(glob.glob(os.path.join(args.fits_path, '*Bp_err.fits')))
+
+def _load_map(d):
+    br, bt, bp = d['b_r'], d['b_t'], d['b_p']
+    br_err, bt_err, bp_err = d['b_r_err'], d['b_t_err'], d['b_p_err']
+
+    b = np.stack([Map(bp).data, -Map(bt).data, Map(br).data]).T
+    b_err = np.stack([Map(bp_err).data, Map(bt_err).data, Map(br_err).data]).T
+
+    return {'b': b, 'b_err': b_err}
+
+d = [{'b_r': b_r_file, 'b_t': b_t_file, 'b_p': b_p_file, 'b_r_err': b_r_err_file, 'b_t_err': b_t_err_file, 'b_p_err': b_p_err_file}
+     for b_r_file, b_t_file, b_p_file, b_r_err_file, b_t_err_file, b_p_err_file in
+     zip(b_r_files, b_t_files, b_p_files, b_r_err_files, b_t_err_files, b_p_err_files)]
+
+with Pool(processes=8) as p:
+    res = [r for r in tqdm(p.imap(_load_map, d), total=len(d))]
+    b_l = [r['b'] for r in res]
+    b_err_l = [r['b_err'] for r in res]
+
+def _load_map_diff(b, b_err, nf2_file):
+    out = CartesianOutput(nf2_file)
+    res = out.load_slice(batch_size=2**15, compute_jacobian=False, metrics=[])
+
+    b_nf2 = res['b'].to_value(u.G)
+
+    b_diff = np.abs(b - b_nf2)
+    b_diff_err = np.clip(b_diff - b_err, a_min=0, a_max=None)
+
+    b_diff = np.linalg.norm(b_diff, axis=-1).mean()
+    b_diff_err = np.linalg.norm(b_diff_err, axis=-1).mean()
+
+    # log_image(b, b_nf2, nf2_file)
+
+    return {'b_diff': b_diff, 'b_diff_err': b_diff_err}
+
+
+def log_image(b, b_nf2, nf2_file):
+    fig, axs = plt.subplots(3, 3, figsize=(10, 10))
+    ax = axs[0, 0]
+    ax.imshow(b[..., 0].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax.set_title('Original')
+    ax = axs[0, 1]
+    ax.imshow(b_nf2[..., 0].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax.set_title('NF2')
+    ax = axs[0, 2]
+    ax.imshow(np.abs(b[..., 0] - b_nf2[..., 0]).T, origin='lower', vmin=0, vmax=1000)
+    ax.set_title('Difference')
+    ax = axs[1, 0]
+    ax.imshow(b[..., 1].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax = axs[1, 1]
+    ax.imshow(b_nf2[..., 1].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax = axs[1, 2]
+    ax.imshow(np.abs(b[..., 1] - b_nf2[..., 1]).T, origin='lower', vmin=0, vmax=1000)
+    ax = axs[2, 0]
+    ax.imshow(b[..., 2].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax = axs[2, 1]
+    ax.imshow(b_nf2[..., 2].T, origin='lower', cmap='gray', vmin=-1000, vmax=1000)
+    ax = axs[2, 2]
+    ax.imshow(np.abs(b[..., 0] - b_nf2[..., 0]).T, origin='lower', vmin=0, vmax=1000)
+    plt.tight_layout()
+    fig.savefig(os.path.join(result_path, os.path.basename(nf2_file) + '.jpg'), dpi=300)
+    plt.close('all')
+    print('diff', np.linalg.norm(b - b_nf2, axis=-1).mean())
+
+
+nf2_files = [sorted(glob.glob(d)) for d in args.nf2_dir]
+nf2_files = [f for d in nf2_files for f in d]
+assert len(b_l) == len(b_err_l) == len(nf2_files)
+b_diff = []
+b_diff_err = []
+for b, b_err, nf2_file in tqdm(zip(b_l, b_err_l, nf2_files), total=len(b_l)):
+    res = _load_map_diff(b, b_err, nf2_file)
+    b_diff.append(res['b_diff'])
+    b_diff_err.append(res['b_diff_err'])
+
+output = load_results(args.pkl_path)
+
+times = output['times']
+metrics = output['metrics']
+height_distribution = output['height_distribution']
+
+Mm_per_pixel = output['Mm_per_pixel']
+
+date_format = DateFormatter('%d-%H:%M')
+
+fig, axs = plt.subplots(3, 1, figsize=(5, 5))
+
+# 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, b_diff, label=r'$L_{\text{B0}}$')
+ax.plot(times, b_diff_err, label=r'$L_{\text{B0, err}}$')
+ax.set_ylabel('$\Delta$B [G]')
+ax.legend()
+
+ax = axs[1]
+ax.plot(times, metrics['divergence'] * 1e4 / Mm_per_pixel)
+ax.set_ylabel(r'$\langle |\nabla \cdot \mathbf{B}$| / |B| $\rangle$' + '\n[$10^{-4}$ Mm$^{-1}$]')
+
+ax = axs[2]
+ax.plot(times, metrics['theta'])
+ax.set_ylabel('$\\theta_{\\text{J}}$ \n[deg]')
+ax.set_xlabel('Time')
+
+ax = axs[2].twinx()
+ax.plot(times, np.sin(np.deg2rad(metrics['theta'])))
+ax.set_ylabel(r'$\sigma_{\text{J}}$')
+
+
+[ax.axvline(x=datetime(2024, 5, 7, 0, 0), linestyle='dotted', c='black') for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(result_path, 'quality.jpg'), dpi=300)
+plt.close()
diff --git a/nf2/evaluation/sharp/plot_separatrix.py b/nf2/evaluation/sharp/plot_separatrix.py
new file mode 100644
index 0000000..e2be001
--- /dev/null
+++ b/nf2/evaluation/sharp/plot_separatrix.py
@@ -0,0 +1,107 @@
+import argparse
+import pickle
+from copy import copy
+
+import h5py
+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.map import Map
+
+from nf2.evaluation.output_metrics import squashing_factor
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Convert NF2 file to VTK.')
+parser.add_argument('--hdf5_file', type=str, help='path to the hdf5 file with the B field.')
+parser.add_argument('--pkl_file', type=str, help='path to the pkl file.')
+parser.add_argument('--euv_file', type=str, help='path to the EUV file.')
+parser.add_argument('--sharp_file', type=str, help='path to the SHARP file.')
+parser.add_argument('--result_path', type=str, help='path to the output directory', required=False, default=None)
+args = parser.parse_args()
+
+result_file = args.result_path
+hdf5_file = args.hdf5_file
+pkl_file = args.pkl_file
+euv_file = args.euv_file
+sharp_file = args.sharp_file
+
+# current density map
+with open(pkl_file, 'rb') as file:
+    output = pickle.load(file)
+current_density_map = output['maps']['current_density_map']
+wcs = output['info']['wcs']
+Mm_per_pixel = output['info']['Mm_per_pixel']
+current_density_map = current_density_map.to_value(u.G * u.cm * u.s ** -1)
+map_wcs = wcs[0]
+
+# squashing factor
+f = h5py.File(hdf5_file, 'r')
+b = f["B"]
+res = squashing_factor(b)
+q_cube = res['q']
+
+# EUV map
+euv_map = Map(euv_file)
+exposure = euv_map.exposure_time.to_value(u.s)
+euv_map = euv_map.reproject_to(map_wcs)
+
+# SHARP map
+sharp_map = Map(sharp_file)
+
+j_norm = LogNorm()
+cm = copy(plt.get_cmap('sdoaia131'))
+cm.set_bad('black')
+
+fig, axs = plt.subplots(4, 1, figsize=(6, 6))
+#
+extent = np.array([0, current_density_map.shape[0],
+                   0, current_density_map.shape[1]]) * Mm_per_pixel
+
+ax = axs[0]
+sharp_im = ax.imshow(sharp_map.data, origin='lower', cmap='gray',
+                     extent=extent, vmin=-3000, vmax=3000)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="5%", pad=0.05)
+fig.colorbar(sharp_im, cax=cax, label='B$_z$ [G]')
+# plot rectangle
+ax.plot([100, 300, 300, 100, 100], [140, 140, 40, 40, 140], 'w--')
+
+#
+ax = axs[1]
+euv_im = ax.imshow(euv_map.data / exposure, origin='lower',
+                   cmap=cm, extent=extent, norm=LogNorm(vmin=5, vmax=1e3))
+ax.contour(sharp_map.data, levels=[-1000, 1000], colors=['black', 'white'], extent=extent, alpha=0.9, linewidths=.7)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="5%", pad=0.05)
+fig.colorbar(euv_im, cax=cax, label='SDO/AIA 131 $\AA$ [DN / s]')
+
+#
+ax = axs[2]
+cd_im = ax.imshow(current_density_map.T, origin='lower',
+                  cmap='inferno', extent=extent, norm=j_norm)
+ax.contour(sharp_map.data, levels=[-1000, 1000], colors=['black', 'white'], extent=extent, alpha=0.9, linewidths=.7)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="5%", pad=0.05)
+fig.colorbar(cd_im, cax=cax, label='J [G cm / s]')
+#
+
+ax = axs[3]
+height = int(5 // Mm_per_pixel)
+q_map = q_cube[:, :, height]
+q_map = np.nan_to_num(q_map, nan=1e3)
+q_im = ax.imshow(q_map.T, origin='lower', cmap='gray',
+                 extent=extent, norm=LogNorm(vmin=1, vmax=1e3))
+divider = make_axes_locatable(ax)
+cax = divider.append_axes("right", size="5%", pad=0.05)
+fig.colorbar(q_im, cax=cax, label='Q')
+
+# axs[-1].set_xlabel('X [Mm]')
+# [ax.set_ylabel('Y [Mm]') for ax in np.ravel(axs)]
+[ax.set_xlim(100, 300) for ax in axs[1:]]
+[ax.set_ylim(40, 140) for ax in axs[1:]]
+
+fig.tight_layout(pad=0.1)
+fig.savefig(result_file, dpi=300)
+plt.close()
diff --git a/nf2/evaluation/sharp/video.py b/nf2/evaluation/sharp/video.py
new file mode 100644
index 0000000..1844dc8
--- /dev/null
+++ b/nf2/evaluation/sharp/video.py
@@ -0,0 +1,104 @@
+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 mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map
+from tqdm import tqdm
+
+from nf2.evaluation.sharp.convert_series import load_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])
+
+    os.makedirs(result_path, exist_ok=True)
+
+    j_norm = LogNorm()
+    free_energy_norm = LogNorm(1e9, 1e14)
+
+    for i, time in tqdm(enumerate(times), total=len(times)):
+        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]
+        free_energy_map[free_energy_map < 1e9] = 1e9
+        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(result_path, f'{time.isoformat("T", timespec="minutes")}.jpg'), dpi=300)
+        plt.close()
+
+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)
+
+    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']
+
+    _plot_map_videos(times, maps, wcs, args.euv_path, result_path, Mm_per_pixel)
+
+if __name__ == '__main__':
+    main()
\ No newline at end of file
diff --git a/nf2/evaluation/spherical/b_diff.py b/nf2/evaluation/spherical/b_diff.py
index d4d2366..a739ce5 100644
--- a/nf2/evaluation/spherical/b_diff.py
+++ b/nf2/evaluation/spherical/b_diff.py
@@ -4,18 +4,21 @@
 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
+from nf2.evaluation.output import SphericalOutput
 
-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'
+nf2_file = '/glade/work/rjarolim/nf2/spherical/2173_full_15_v02/extrapolation_result.nf2'
+results_path = '/glade/work/rjarolim/nf2/spherical/2173_full_15_v02/results'
+ref_r_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits'
+ref_t_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt.fits'
+ref_p_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp.fits'
+ref_r_err_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br_err.fits'
+ref_t_err_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bt_err.fits'
+ref_p_err_file = '/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Bp_err.fits'
 os.makedirs(results_path, exist_ok=True)
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
@@ -23,39 +26,61 @@
 ref_r_map = Map(ref_r_file)
 ref_t_map = Map(ref_t_file)
 ref_p_map = Map(ref_p_file)
+
+ref_r_err_map = Map(ref_r_err_file)
+ref_t_err_map = Map(ref_t_err_file)
+ref_p_err_map = Map(ref_p_err_file)
+
+err_b = np.stack([ref_r_err_map.data, ref_t_err_map.data, ref_p_err_map.data]).T
+
 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
+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,
+    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'])
+model = SphericalOutput(nf2_file)
+model_out = model.load_coords(cartesian_coords)
 
-b = load_coords(model, 1, state['b_norm'], cartesian_coords, device, progress=True, compute_currents=False)
+b_xyz = model_out['b']
+b = vector_cartesian_to_spherical(b_xyz, spherical_coords).to_value(u.G)
 
-b = vector_cartesian_to_spherical(b, spherical_coords)
+b_diff = np.clip(np.abs(b - ref_b) - err_b, a_min=0, a_max=None)
 
 # create difference plot
-fig, axs = plt.subplots(3, 3, figsize=(15, 15))
+fig, axs = plt.subplots(3, 3, figsize=(15, 10))
 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')
+    ax = axs[i, 0]
+    im = ax.imshow(ref_b[..., i].T, cmap='grey', vmin=-500, vmax=500)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='[G]')
+
+    ax = axs[i, 1]
+    im = ax.imshow(b[..., i].T, cmap='grey', vmin=-500, vmax=500)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='[G]')
+
+    ax = axs[i, 2]
+    im = ax.imshow(b_diff[..., i].T, cmap='viridis', vmin=0, vmax=100)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+    plt.colorbar(im, cax=cax, label='[G]')
+
+axs[0, 0].set_title('Reference')
+axs[0, 1].set_title('Model')
+axs[0, 2].set_title('Difference')
+[ax.set_axis_off() for ax in axs.ravel()]
+
 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))
+print('Mean absolute difference:', np.nanmean(np.linalg.norm(b_diff, axis=-1)))
diff --git a/nf2/evaluation/spherical/compare.py b/nf2/evaluation/spherical/compare.py
index ffc4745..598bf33 100644
--- a/nf2/evaluation/spherical/compare.py
+++ b/nf2/evaluation/spherical/compare.py
@@ -8,8 +8,8 @@
 
 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.evaluation.output import CartesianOutput, SphericalOutput
+from nf2.evaluation.output_metrics import current_density, spherical_energy_gradient, energy_gradient
 from nf2.potential.potential_field import get_potential_field
 
 Mm_per_pixel = 0.72
diff --git a/nf2/evaluation/spherical/currents.py b/nf2/evaluation/spherical/currents.py
index ade32c0..a19c96e 100644
--- a/nf2/evaluation/spherical/currents.py
+++ b/nf2/evaluation/spherical/currents.py
@@ -1,146 +1,286 @@
+import argparse
 import glob
 import os
 
 import numpy as np
 import pfsspy
 import torch
+from astropy import units as u
 from astropy.coordinates import SkyCoord
+from dateutil.parser import parse
 from matplotlib import pyplot as plt
-from matplotlib.colors import LogNorm, SymLogNorm
+from matplotlib.colors import LogNorm, Normalize
 from mpl_toolkits.axes_grid1 import make_axes_locatable
-from sunpy.map import Map
-from torch import nn
+from sunpy.map import Map, all_coordinates_from_map, make_heliographic_header
 
-from nf2.data.util import spherical_to_cartesian, vector_spherical_to_cartesian
-from nf2.evaluation.unpack import load_coords
+from nf2.data.util import vector_cartesian_to_spherical
+from nf2.evaluation.metric import energy
+from nf2.evaluation.output import SphericalOutput
 
-from astropy import units as u
+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('--out_path', type=str, help='path to the target VTK files', required=False, default=None)
+
+    args = parser.parse_args()
+
+    files = args.nf2_path
+    results_path = args.out_path if args.out_path is not None else os.path.dirname(args.nf2_path)
+
+    synoptic_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"
+    }
+    full_disc_files = {
+        'Br': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Br.fits")),
+        'Bt': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Bt.fits")),
+        'Bp': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Bp.fits"))
+    }
+    euv_files = sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/euv/*.193.image_lev1.fits"))
+
+    euv_dates = np.array([parse(os.path.basename(f).split('.')[2]) for f in euv_files])
+    mag_dates = np.array([parse(os.path.basename(f).split('.')[2].replace('_TAI', 'Z').replace('_', 'T')) for f in full_disc_files['Br']])
+
+    os.makedirs(results_path, exist_ok=True)
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    synoptic_map_br = Map(synoptic_files['Br'])
+    synoptic_map_bt = Map(synoptic_files['Bt'])
+    synoptic_map_bp = Map(synoptic_files['Bp'])
+
+
+    def _build_synoptic_map(br_file):
+        br = Map(br_file)
+        synoptic_map_br.meta['date-obs'] = br.date.to_datetime().date_str()
+        br = br.reproject_to(synoptic_map_br.wcs)
+
+        br_data = synoptic_map_br.data
+        br_data[~np.isnan(br.data)] = br.data[~np.isnan(br.data)]
+
+        br_map = Map(br_data, synoptic_map_br.meta)
+        return br_map
+
+
+    sampling = [128, 128, 128]
+    longitude_range = [150, 210]
+    latitude_range = [60, 120]
+    radius_range = [1.0, 1.3]
+    dr = ((radius_range[1] - radius_range[0]) * u.solRad).to_value(u.cm) / sampling[0]
+
+    alpha_norm = Normalize(vmin=0, vmax=10)
+    j_norm = LogNorm(vmin=1e10, vmax=3e12)
+    energy_norm = LogNorm(vmin=1e10, vmax=5e13)
+    euv_norm = LogNorm(vmin=50, vmax=5e3)
+
+    files = sorted(glob.glob(files, recursive=True))
+
+    for file in files:
+        print(file)
+        # date = parse(os.path.basename(file).split('.')[0].replace('_TAI', 'Z').replace('_', 'T'))
+        ds = file.split(os.sep)[-2].split('_')
+        date = parse(f'{ds[0]}T{ds[1]}Z')
+        date_str = date.isoformat("T", timespec="minutes").replace('+00:00', '')
+        model = SphericalOutput(file)
+        # prepare synoptic data
+        # br_map = _build_synoptic_map(full_disc_files['Br'][i])
+        # br_map = br_map.resample([360, 180] * u.pix)
+
+        br_map = synoptic_map_br.resample([360, 180] * u.pix)
+        spherical_coords = all_coordinates_from_map(br_map)
+        model_out = model.load_spherical_coords(spherical_coords)
+        b_r = vector_cartesian_to_spherical(model_out['b'].to_value(u.G),
+                                            model_out['spherical_coords'])[..., 0]
+        br_map.data[:] = b_r
+
+        # build PFSS model
+        pfss_in = pfsspy.Input(br_map, 100, 2.5)
+        pfss_out = pfsspy.pfss(pfss_in)
+
+        # EUV
+        euv_file = euv_files[np.argmin(np.abs(euv_dates - date))]
+        euv_map = Map(euv_file)
+
+        exposure_time = euv_map.exposure_time.to_value(u.s)
+        carr_header = make_heliographic_header(euv_map.date, euv_map.observer_coordinate, (180 * 8, 360 * 8),
+                                               frame='carrington', projection_code='CAR',
+                                               map_center_longitude=180 * u.deg)
+        euv_map = euv_map.reproject_to(carr_header)
+
+        # MAG
+        mag_file = full_disc_files['Br'][np.argmin(np.abs(mag_dates - date))]
+        mag_map = Map(mag_file)
+        carr_header = make_heliographic_header(mag_map.date, mag_map.observer_coordinate, (180 * 8, 360 * 8),
+                                               frame='carrington', projection_code='CAR',
+                                               map_center_longitude=180 * u.deg)
+        mag_map = mag_map.reproject_to(carr_header)
+
+        fig, ax = plt.subplots(figsize=(5, 5))
+
+        ax.imshow(euv_map.data / exposure_time, origin='lower', cmap=euv_map.cmap, extent=[0, 360, 180, 0],
+                  norm=euv_norm)
+        ax.set_xlim(*longitude_range)
+        ax.set_ylim(*reversed(latitude_range))
+        fig.savefig(os.path.join(results_path, f'euv_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        model_out = model.load_spherical(radius_range * u.solRad,
+                                         longitude_range=longitude_range * u.deg, latitude_range=latitude_range * u.deg,
+                                         metrics=['j', 'alpha'], sampling=sampling, progress=True)
+
+        spherical_coords = model_out['spherical_coords']
+        spherical_coords = SkyCoord(lon=spherical_coords[..., 2] * u.rad,
+                                    lat=(np.pi / 2 - spherical_coords[..., 1]) * u.rad,
+                                    radius=spherical_coords[..., 0] * u.solRad,
+                                    frame=br_map.coordinate_frame)
+        pfss_cube_shape = spherical_coords.shape
+        pfss_b = pfss_out.get_bvec(spherical_coords.reshape((-1,)))
+        pfss_b = pfss_b.reshape(*pfss_cube_shape, 3)
+        potential_energy = energy(pfss_b)
+
+        j = model_out['metrics']['j'].to_value(u.G / u.s)
+        b = model_out['b'].to_value(u.G)
+        alpha = model_out['metrics']['alpha'].to_value(u.cm ** -1)
+        ff_energy = energy(b)
+        free_energy = ff_energy - potential_energy
+
+        extent = [*longitude_range, *reversed(latitude_range)]
+
+        # B field
+        fig, ax = plt.subplots(figsize=(5, 5))
+        b_r = vector_cartesian_to_spherical(b, model_out['spherical_coords'])[0, ..., 0]
+        im = ax.imshow(b_r, origin='upper', vmin=-500, vmax=500,
+                       cmap='gray', extent=extent)
+        ax.set_title('B field')
+        ax.set_xlabel('Longitude [deg]')
+        ax.set_ylabel('Latitude [deg]')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical')
+        fig.savefig(os.path.join(results_path, f'b_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        # B observed
+        fig, ax = plt.subplots(figsize=(5, 5))
+        b_r = mag_map.data
+        im = ax.imshow(b_r, origin='lower', vmin=-500, vmax=500,
+                       cmap='gray', extent=[0, 360, 180, 0])
+        ax.set_title('B observed')
+        ax.set_xlabel('Longitude [deg]')
+        ax.set_ylabel('Latitude [deg]')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical')
+        ax.set_xlim(*longitude_range)
+        ax.set_ylim(*reversed(latitude_range))
+        fig.savefig(os.path.join(results_path, f'b_obs_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+
+
+        # currents
+        fig, ax = plt.subplots(figsize=(5, 5))
+        current_density = np.linalg.norm(j, axis=-1)
+        im = ax.imshow(current_density.sum(0) * dr, origin='upper', norm=j_norm, cmap='inferno', extent=extent)
+        ax.set_title('Currents')
+        ax.set_xlabel('Longitude [deg]')
+        ax.set_ylabel('Latitude [deg]')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical')
+        fig.savefig(os.path.join(results_path, f'currents_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        alpha[np.linalg.norm(b, axis=-1) < 5] = 0
+        # alpha
+        fig, ax = plt.subplots(figsize=(5, 5))
+        im = ax.imshow(alpha.sum(0) * dr, origin='upper', cmap='viridis', extent=extent, norm=alpha_norm)
+        ax.set_title(r'$\alpha$')
+        ax.set_xlabel('Longitude [deg]')
+        ax.set_ylabel('Latitude [deg]')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical')
+        fig.savefig(os.path.join(results_path, f'alpha_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+
+        # free energy
+        def _plot_energy(ax, energy, title):
+            im = ax.imshow(energy, origin='upper', cmap='jet', extent=extent, norm=energy_norm)
+            ax.set_title(title)
+            ax.set_xlabel('Longitude [deg]')
+            ax.set_ylabel('Latitude [deg]')
+            divider = make_axes_locatable(ax)
+            cax = divider.append_axes('right', size='2%', pad=0.05)
+            fig.colorbar(im, cax=cax, orientation='vertical', label='[erg/cm^2]')
+
+
+        integrated_free_energy = free_energy.sum(0) * dr
+        integrated_free_energy[integrated_free_energy < 1e8] = 1e8
+        integrated_potential_energy = potential_energy.sum(0) * dr
+        integrated_ff_energy = ff_energy.sum(0) * dr
+
+        fig, ax = plt.subplots(figsize=(10, 5))
+        _plot_energy(ax, integrated_free_energy, 'Free Energy')
+        fig.savefig(os.path.join(results_path, f'free_energy_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        fig, ax = plt.subplots(figsize=(5, 5))
+        _plot_energy(ax, integrated_potential_energy, 'Potential Energy')
+        fig.savefig(os.path.join(results_path, f'potential_energy_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        fig, ax = plt.subplots(figsize=(5, 5))
+        _plot_energy(ax, integrated_ff_energy, 'Force-Free Energy')
+        fig.savefig(os.path.join(results_path, f'ff_energy_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
+
+        # plot overview video
+        fig, axs = plt.subplots(1, 4, figsize=(20, 5))
+
+        ax = axs[0]
+        im = ax.imshow(euv_map.data / exposure_time, origin='lower', cmap=euv_map.cmap, extent=[0, 360, 0, 180],
+                       norm=euv_norm)
+        ax.set_title('SDO/AIA 193 $\AA$')
+        ax.set_xlim(*longitude_range)
+        ax.set_ylim(*latitude_range)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical', label='[DN/s]')
+
+        # ax = axs[1]
+        # im = ax.imshow(b_r, origin='upper', vmin=-500, vmax=500, cmap='gray', extent=extent)
+        # ax.set_title('$B_{z}$')
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes('right', size='2%', pad=0.05)
+        # fig.colorbar(im, cax=cax, orientation='vertical', label='[G]')
+
+        ax = axs[1]
+        im = ax.imshow(current_density.sum(0) * dr, origin='upper', norm=j_norm, cmap='inferno', extent=extent)
+        ax.set_title('Current density')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical', label='[G cm / s]')
+
+        ax = axs[2]
+        im = ax.imshow(alpha.sum(0) * dr, origin='upper', cmap='viridis', extent=extent, norm=alpha_norm)
+        ax.set_title(r'$\alpha$')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical', label='[cm$^{-1}$]')
 
-# 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)
+        ax = axs[3]
+        im = ax.imshow(integrated_free_energy, origin='upper', cmap='jet', extent=extent, norm=energy_norm)
+        ax.set_title('Free Energy')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes('right', size='2%', pad=0.05)
+        fig.colorbar(im, cax=cax, orientation='vertical', label='[erg/cm$^2$]')
 
+        axs[0].set_ylabel('Latitude [deg]')
+        [ax.set_xlabel('Longitude [deg]') for ax in axs]
 
+        fig.tight_layout()
 
+        fig.savefig(os.path.join(results_path, f'video_{date_str}.jpg'), dpi=300)
+        plt.close(fig)
diff --git a/nf2/evaluation/spherical/disk.py b/nf2/evaluation/spherical/disk.py
new file mode 100644
index 0000000..d6e9999
--- /dev/null
+++ b/nf2/evaluation/spherical/disk.py
@@ -0,0 +1,159 @@
+import argparse
+import glob
+import os
+
+import numpy as np
+import torch
+from astropy import units as u
+from astropy.coordinates import SkyCoord
+from dateutil.parser import parse
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, Normalize
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+from sunpy.map import Map, all_coordinates_from_map
+
+from nf2.data.util import vector_cartesian_to_spherical
+from nf2.evaluation.metric import energy
+from nf2.evaluation.output import SphericalOutput
+
+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('--out_path', type=str, help='path to the target VTK files', required=False, default=None)
+
+    args = parser.parse_args()
+
+    files = args.nf2_path
+    results_path = args.out_path if args.out_path is not None else os.path.dirname(args.nf2_path)
+
+    synoptic_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"
+    }
+    full_disc_files = {
+        'Br': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Br.fits")),
+        'Bt': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Bt.fits")),
+        'Bp': sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/full_disk/*Bp.fits"))
+    }
+    euv_files = sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/euv/*.193.image_lev1.fits"))
+
+    euv_dates = np.array([parse(os.path.basename(f).split('.')[2]) for f in euv_files])
+    mag_dates = np.array([parse(os.path.basename(f).split('.')[2].replace('_TAI', 'Z').replace('_', 'T')) for f in
+                          full_disc_files['Br']])
+
+    os.makedirs(results_path, exist_ok=True)
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    overwrite = False
+    img_res = (256, 256)
+    sampling = 128
+    radius_range = [1.0, 1.3]
+    dr = ((radius_range[1] - radius_range[0]) * u.solRad).to(u.cm) / sampling
+
+    alpha_norm = Normalize(vmin=0, vmax=10)
+    j_norm = LogNorm(vmin=1e10, vmax=3e12)
+    energy_norm = LogNorm(vmin=1e10, vmax=5e13)
+    euv_norm = LogNorm(vmin=50, vmax=5e3)
+
+    files = sorted(glob.glob(files))
+
+    for i, file in enumerate(files):
+        print(f'Processing {file} ({i + 1}/{len(files)})')
+        date = parse(os.path.basename(file).split('.')[0].replace('_TAI', 'Z').replace('_', 'T'))
+
+        img_path = os.path.join(results_path, f'disk_{date.strftime("%Y%m%dT%H%M%S")}.png')
+
+        if os.path.exists(img_path) and not overwrite:
+            continue
+
+        model = SphericalOutput(file)
+
+        # EUV
+        euv_file = euv_files[np.argmin(np.abs(euv_dates - date))]
+        euv_map = Map(euv_file)
+
+        exposure_time = euv_map.exposure_time.to_value(u.s)
+        euv_map = euv_map.resample(img_res * u.pix)
+        euv_coords = all_coordinates_from_map(euv_map)
+        radius = np.sqrt(euv_coords.Tx ** 2 + euv_coords.Ty ** 2) / euv_map.rsun_obs
+        alpha_mask = np.clip(1.3 - radius, 0, None) / 0.3
+        alpha_mask = alpha_mask ** 3
+        alpha_mask[radius < 1.0] = 1
+
+        # mag
+        # mag_file = full_disc_files['Br'][np.argmin(np.abs(mag_dates - date))]
+        # mag_map = Map(mag_file)
+        # mag_map = mag_map.reproject_to(euv_map.wcs)  # transform to smaller (HMI) subframe
+
+        # load coordinates
+        coords = all_coordinates_from_map(euv_map).transform_to('heliographic_carrington')
+        coords = np.stack([
+            np.ones_like(coords.lat.to_value(u.rad)),
+            coords.lat.to_value(u.rad),
+            coords.lon.to_value(u.rad)], axis=-1)
+        coords = coords[None].repeat(sampling, axis=0)
+        radius_array = np.linspace(radius_range[0], radius_range[1], sampling)
+        coords[..., 0] = radius_array[:, None, None]
+
+        skycoords = SkyCoord(lon=coords[..., 2] * u.rad, lat=coords[..., 1] * u.rad, radius=coords[..., 0] * u.solRad,
+                             frame='heliographic_carrington')
+
+        model_out = model.load_spherical_coords(skycoords, metrics=['j', 'alpha'], progress=True,
+                                                batch_size=int(2 ** 14))
+
+        br = vector_cartesian_to_spherical(model_out['b'], model_out['spherical_coords'])[0, ..., 0]
+
+        j = np.linalg.norm(model_out['metrics']['j'], axis=-1).sum(0) * dr
+        j = j.to_value(u.G * u.cm / u.s)
+
+        ff_energy = energy(model_out['b']).sum(0) * dr
+        ff_energy = ff_energy.to_value(u.G ** 2 * u.cm)
+
+        fig, axs = plt.subplots(1, 3, figsize=(12, 4))
+
+        ax = axs[0]
+        im = ax.imshow(euv_map.data / exposure_time, norm=euv_norm, cmap=euv_map.cmap, origin='lower', alpha=alpha_mask)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="2%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Intensity [DN/s]', orientation='vertical')
+
+        ax = axs[1]
+        im = ax.imshow(j, norm=j_norm, cmap='inferno', origin='lower')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="2%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Current Density [G cm / s]', orientation='vertical')
+
+        ax = axs[2]
+        im = ax.imshow(ff_energy, norm=energy_norm, cmap='jet', origin='lower')
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="2%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Magnetic Energy [erg / cm^2]', orientation='vertical')
+
+        [ax.set_axis_off() for ax in axs]
+        axs[0].set_title(date.strftime('%Y-%m-%d %H:%M:%S'))
+
+        fig.tight_layout()
+        fig.savefig(img_path, dpi=300)
+        plt.close(fig)
+
+        # fig, axs = plt.subplots(1, 2, figsize=(8, 4))
+        #
+        # ax = axs[0]
+        # im = ax.imshow(mag_map.data, cmap='gray', origin='lower', vmin=-500, vmax=500)
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="2%", pad=0.05)
+        # plt.colorbar(im, cax=cax, label='$B_\\text{r}$ [G]', orientation='vertical')
+        #
+        # ax = axs[1]
+        # im = ax.imshow(br.to_value(u.G), cmap='gray', origin='lower', vmin=-500, vmax=500)
+        # divider = make_axes_locatable(ax)
+        # cax = divider.append_axes("right", size="2%", pad=0.05)
+        # plt.colorbar(im, cax=cax, label='$B_\\text{r}$ [G]', orientation='vertical')
+        #
+        # [ax.set_axis_off() for ax in axs]
+        # axs[0].set_title(date.strftime('%Y-%m-%d %H:%M:%S'))
+        # fig.tight_layout()
+        # mag_img_path = os.path.join(results_path, f'mag_{date.strftime("%Y%m%dT%H%M%S")}.png')
+        # fig.savefig(mag_img_path, dpi=300)
+        # plt.close(fig)
diff --git a/nf2/evaluation/spherical/energy_diff.py b/nf2/evaluation/spherical/energy_diff.py
new file mode 100644
index 0000000..b81027d
--- /dev/null
+++ b/nf2/evaluation/spherical/energy_diff.py
@@ -0,0 +1,66 @@
+import glob
+import os
+
+import numpy as np
+from astropy import units as u
+from dateutil.parser import parse
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, Normalize, SymLogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.evaluation.metric import energy
+from nf2.evaluation.output import SphericalOutput
+
+file_1 = f'/glade/work/rjarolim/nf2/spherical/2173/20160205_200000_v02/extrapolation_result.nf2'
+file_2 = f'/glade/work/rjarolim/nf2/spherical/2173/20160205_220000_v02/extrapolation_result.nf2'
+results_path = '/glade/work/rjarolim/nf2/spherical/2173/results'
+
+euv_files = sorted(glob.glob("/glade/work/rjarolim/data/global/fd_2173/euv/*.193.image_lev1.fits"))
+euv_dates = np.array([parse(os.path.basename(f).split('.')[2]) for f in euv_files])
+
+os.makedirs(results_path, exist_ok=True)
+
+sampling = [128, 128, 128]
+longitude_range = [150, 210]
+latitude_range = [60, 120]
+radius_range = [1.0, 1.3]
+dr = ((radius_range[1] - radius_range[0]) * u.solRad).to_value(u.cm) / sampling[0]
+
+alpha_norm = Normalize(vmin=0, vmax=10)
+j_norm = LogNorm(vmin=1e10, vmax=3e12)
+energy_norm = LogNorm(vmin=1e10, vmax=5e13)
+euv_norm = LogNorm(vmin=50, vmax=5e3)
+
+model_1 = SphericalOutput(file_1)
+model_out_1 = model_1.load_spherical(radius_range * u.solRad,
+                                     longitude_range=longitude_range * u.deg, latitude_range=latitude_range * u.deg,
+                                     metrics=['j', 'alpha'], sampling=sampling, progress=True)
+
+model_2 = SphericalOutput(file_2)
+model_out_2 = model_2.load_spherical(radius_range * u.solRad,
+                                     longitude_range=longitude_range * u.deg, latitude_range=latitude_range * u.deg,
+                                     metrics=['j', 'alpha'], sampling=sampling, progress=True)
+
+energy_1 = energy(model_out_1['b'].to_value(u.G))
+energy_2 = energy(model_out_2['b'].to_value(u.G))
+
+# energy_diff = -np.clip(energy_1 - energy_2, a_min=None, a_max=0)
+# energy_diff = energy_2 - energy_1
+energy_diff = np.clip(energy_2 - energy_1, a_min=0, a_max=None)
+
+energy_map = energy_diff.sum(0) * dr
+
+extent = [*longitude_range, *reversed(latitude_range)]
+
+fig, ax = plt.subplots(figsize=(5, 5))
+im = ax.imshow(energy_map, origin='upper', cmap='Reds', extent=extent,
+               # vmax=1e11, vmin=-1e11)
+               norm=LogNorm(vmin=1e10, vmax=1e12))
+               # norm=SymLogNorm(vmin=-1e12, vmax=1e12, linthresh=1e10))
+ax.set_xlabel('Longitude [deg]')
+ax.set_ylabel('Latitude [deg]')
+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, 'energy_diff.jpg'), dpi=300)
+plt.close(fig)
diff --git a/nf2/evaluation/spherical/input_plot.py b/nf2/evaluation/spherical/input_plot.py
new file mode 100644
index 0000000..aa4aadc
--- /dev/null
+++ b/nf2/evaluation/spherical/input_plot.py
@@ -0,0 +1,20 @@
+import numpy as np
+from matplotlib import pyplot as plt, cm
+from matplotlib.colors import Normalize
+from sunpy.map import Map
+
+synoptic_map = Map('/glade/work/rjarolim/data/global/fd_2173/hmi.synoptic_mr_polfil_720s.2173.Mr_polfil.fits')
+br_map = Map('/glade/work/rjarolim/data/global/fd_2173/full_disk/hmi.b_720s.20160205_000000_TAI.Br.fits')
+
+img = Normalize(vmin=-500, vmax=500)(br_map.data)
+img = cm.get_cmap('gray')(img)
+plt.imsave('/glade/work/rjarolim/nf2/spherical/br.png', img)
+
+synoptic_map.meta['date-obs'] = br_map.date.to_datetime().isoformat()
+reprojected = br_map.reproject_to(synoptic_map.wcs)
+synoptic_map.data[~np.isnan(reprojected.data)] = np.nan
+
+img = Normalize(vmin=-500, vmax=500)(synoptic_map.data)
+
+img = cm.get_cmap('gray')(img)
+plt.imsave('/glade/work/rjarolim/nf2/spherical/synoptic.png', img)
diff --git a/nf2/evaluation/sst/compare_disambiguation.py b/nf2/evaluation/sst/compare_disambiguation.py
deleted file mode 100644
index 8138dc3..0000000
--- a/nf2/evaluation/sst/compare_disambiguation.py
+++ /dev/null
@@ -1,75 +0,0 @@
-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/sst/compare_muram_mfr.py b/nf2/evaluation/sst/compare_muram_mfr.py
new file mode 100644
index 0000000..9b4dacc
--- /dev/null
+++ b/nf2/evaluation/sst/compare_muram_mfr.py
@@ -0,0 +1,180 @@
+import argparse
+import os
+
+import numpy as np
+from astropy import units as u, constants
+from astropy.io import fits
+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 b_nabla_bz, curl
+from nf2.evaluation.output import CartesianOutput
+from nf2.loader.muram import MURaMSnapshot
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Evaluate SHARP snapshot.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+out_path = '/glade/work/rjarolim/nf2/disambiguation/evaluation'
+os.makedirs(out_path, exist_ok=True)
+
+
+Mm_per_pixel = 0.192
+
+mfr_muram_snapshot = MURaMSnapshot('/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_B/3D', iteration=474000)
+mfr_muram_cube = mfr_muram_snapshot.load_cube(resolution=Mm_per_pixel * u.Mm / u.pix, height=50 * u.Mm)
+mfr_muram_b = mfr_muram_cube['B']
+mfr_muram_BnablaBz = b_nabla_bz(mfr_muram_b) / Mm_per_pixel
+mfr_muram_j = (curl(mfr_muram_b) / Mm_per_pixel * u.G / u.Mm) * constants.c / (4 * np.pi)
+mfr_muram_j = mfr_muram_j.to_value(u.G / u.s)
+
+mfr_1slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_mfr_1slices_v01/extrapolation_result.nf2')
+mfr_2slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_mfr_2slices_v01/extrapolation_result.nf2')
+mfr_2slices_amb_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_mfr_2slices_ambiguous_v01/extrapolation_result.nf2')
+
+x_min, x_max = mfr_1slice_model.coord_range[0]
+y_min, y_max = mfr_1slice_model.coord_range[1]
+
+Mm_per_ds = mfr_1slice_model.Mm_per_ds
+ds_per_pixel = Mm_per_pixel / Mm_per_ds
+
+coords = np.stack(
+    np.meshgrid(np.linspace(x_min, x_max, np.round((x_max - x_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(y_min, y_max, np.round((y_max - y_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(0, 50 / Mm_per_ds, np.round(50 / Mm_per_pixel + 1).astype(int)),
+                indexing='ij'), -1)
+mfr_1slice_out = mfr_1slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+mfr_2slice_out = mfr_2slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+mfr_2slices_amb_out = mfr_2slices_amb_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+
+##############################################################################
+
+extent = np.array([x_min, x_max, y_min, y_max]) * Mm_per_ds
+
+fig, ax = plt.subplots(1, 1, figsize=(10, 5))
+
+im = ax.imshow(mfr_muram_b[:, :, 0, 2].T, origin='lower', cmap='gray',
+                 vmin=-1000, vmax=1000, extent=extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+
+y_slice = 200
+# draw red line at slice position
+ax.axvline(y_slice * Mm_per_pixel, color='red', linestyle='--')
+
+plt.savefig(os.path.join(out_path, 'muram_mfr_bz.png'))
+plt.close()
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, 0, 70]
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(mfr_muram_BnablaBz[y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(mfr_1slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(mfr_2slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(mfr_2slices_amb_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+                 cmap='RdBu_r', vmin=-.1, vmax=.1, extent=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'$\hat{B} \cdot \nabla \hat{B}_z$ [Mm$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('Y [Mm]') for ax in axs]
+axs[0].set_ylabel('Z [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_mfr_b_nabla_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm(vmin=1e-1, vmax=1e3)
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(np.linalg.norm(mfr_muram_b[y_slice, :, :], axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(np.linalg.norm(mfr_1slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(np.linalg.norm(mfr_2slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(np.linalg.norm(mfr_2slices_amb_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\hat{B} \cdot \nabla \hat{B}_z$ [Mm$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_mfr_b.png'), dpi=300, transparent=True)
+plt.close()
+
+
+
+##############################################################################
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, y_min * Mm_per_ds, y_max * Mm_per_ds]
+j_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(np.linalg.norm(mfr_muram_j, axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(np.linalg.norm(mfr_1slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(np.linalg.norm(mfr_2slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(np.linalg.norm(mfr_2slices_amb_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G cm s$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_mfr_j.png'), dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sst/compare_muram_sma.py b/nf2/evaluation/sst/compare_muram_sma.py
new file mode 100644
index 0000000..66db530
--- /dev/null
+++ b/nf2/evaluation/sst/compare_muram_sma.py
@@ -0,0 +1,178 @@
+import argparse
+import os
+
+import numpy as np
+from astropy import units as u, constants
+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 b_nabla_bz, curl
+from nf2.evaluation.output import CartesianOutput
+from nf2.loader.muram import MURaMSnapshot
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Evaluate SHARP snapshot.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+out_path = '/glade/work/rjarolim/nf2/disambiguation/evaluation'
+os.makedirs(out_path, exist_ok=True)
+
+Mm_per_pixel = 0.192
+
+sma_muram_snapshot = MURaMSnapshot('/glade/campaign/hao/radmhd/Rempel/Spot_Motion/case_A/3D', iteration=399000)
+sma_muram_cube = sma_muram_snapshot.load_cube(resolution=Mm_per_pixel * u.Mm / u.pix, height=50 * u.Mm)
+sma_muram_b = sma_muram_cube['B']
+sma_muram_BnablaBz = b_nabla_bz(sma_muram_b) / Mm_per_pixel
+sma_muram_j = (curl(sma_muram_b) / Mm_per_pixel * u.G / u.Mm) * constants.c / (4 * np.pi)
+sma_muram_j = sma_muram_j.to_value(u.G / u.s)
+
+sma_1slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_sma_1slices_v01/extrapolation_result.nf2')
+sma_2slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_sma_2slices_v01/extrapolation_result.nf2')
+sma_2slices_amb_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/disambiguation/muram_sma_2slices_ambiguous_v01/extrapolation_result.nf2')
+
+x_min, x_max = sma_1slice_model.coord_range[0]
+y_min, y_max = sma_1slice_model.coord_range[1]
+
+Mm_per_ds = sma_1slice_model.Mm_per_ds
+ds_per_pixel = Mm_per_pixel / Mm_per_ds
+
+coords = np.stack(
+    np.meshgrid(np.linspace(x_min, x_max, np.round((x_max - x_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(y_min, y_max, np.round((y_max - y_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(0, 50 / Mm_per_ds, np.round(50 / Mm_per_pixel + 1).astype(int)),
+                indexing='ij'), -1)
+sma_1slice_out = sma_1slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+sma_2slice_out = sma_2slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+sma_2slices_amb_out = sma_2slices_amb_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+
+##############################################################################
+
+extent = np.array([x_min, x_max, y_min, y_max]) * Mm_per_ds
+
+fig, ax = plt.subplots(1, 1, figsize=(10, 5))
+
+im = ax.imshow(sma_muram_b[:, :, 0, 2].T, origin='lower', cmap='gray',
+                 vmin=-1000, vmax=1000, extent=extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+
+y_slice = 200
+# draw red line at slice position
+ax.axvline(y_slice * Mm_per_pixel, color='red', linestyle='--')
+
+plt.savefig(os.path.join(out_path, 'muram_sma_bz.png'))
+plt.close()
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, 0, 70]
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(sma_muram_BnablaBz[y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(sma_1slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(sma_2slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(sma_2slices_amb_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+                 cmap='RdBu_r', vmin=-.1, vmax=.1, extent=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'$\hat{B} \cdot \nabla \hat{B}_z$ [Mm$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('Y [Mm]') for ax in axs]
+axs[0].set_ylabel('Z [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_sma_b_nabla_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm(vmin=1e-1, vmax=1e3)
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(np.linalg.norm(sma_muram_b[y_slice, :, :], axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(np.linalg.norm(sma_1slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(np.linalg.norm(sma_2slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(np.linalg.norm(sma_2slices_amb_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\hat{B} \cdot \nabla \hat{B}_z$ [Mm$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_sma_b.png'), dpi=300, transparent=True)
+plt.close()
+
+
+
+##############################################################################
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, y_min * Mm_per_ds, y_max * Mm_per_ds]
+j_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 4, figsize=(15, 2.7))
+
+ax = axs[0]
+im = ax.imshow(np.linalg.norm(sma_muram_j, axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('MURaM')
+
+ax = axs[1]
+ax.imshow(np.linalg.norm(sma_1slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('NF2 - single height')
+
+ax = axs[2]
+ax.imshow(np.linalg.norm(sma_2slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+ax.set_title('NF2 - multi height')
+
+
+ax = axs[3]
+ax.imshow(np.linalg.norm(sma_2slices_amb_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G cm s$^{-1}$]')
+ax.set_title('NF2 - multi height (ambiguous)')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'muram_sma_j.png'), dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sst/compare_sharp_disambiguation.py b/nf2/evaluation/sst/compare_sharp_disambiguation.py
new file mode 100644
index 0000000..54a2875
--- /dev/null
+++ b/nf2/evaluation/sst/compare_sharp_disambiguation.py
@@ -0,0 +1,10 @@
+import argparse
+
+from nf2.evaluation.output import CartesianOutput
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Evaluate SHARP snapshot.')
+    parser.add_argument('--output', type=str, help='output path.')
+    args = parser.parse_args()
+
+    sharp_out = CartesianOutput('/glade/work/rjarolim/nf2/sst/sharp_13392_ambiguous_v02/sharp_13392_ambiguous_v02.vtk')
diff --git a/nf2/evaluation/sst/compare_sst.py b/nf2/evaluation/sst/compare_sst.py
new file mode 100644
index 0000000..3001c6d
--- /dev/null
+++ b/nf2/evaluation/sst/compare_sst.py
@@ -0,0 +1,237 @@
+import argparse
+import os
+
+import numpy as np
+from astropy import units as u, constants
+from astropy.io import fits
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm, Normalize
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.evaluation.metric import b_nabla_bz, curl
+from nf2.evaluation.output import CartesianOutput, HeightTransformOutput
+from nf2.loader.muram import MURaMSnapshot
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Evaluate SST extrapolation.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+out_path = '/glade/work/rjarolim/nf2/sst/evaluation'
+os.makedirs(out_path, exist_ok=True)
+
+# select intensity data in wings
+sst_data = fits.getdata('/glade/work/rjarolim/data/SST/panorama_8542_StkI.fits')[10]
+
+
+mfr_1slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/sst/13392_1slices_0851_v01/extrapolation_result.nf2')
+mfr_2slice_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/sst/13392_2slices_0851_v02/extrapolation_result.nf2')
+
+x_min, x_max = mfr_1slice_model.coord_range[0]
+y_min, y_max = mfr_1slice_model.coord_range[1]
+
+Mm_per_pixel = 0.36
+Mm_per_ds = mfr_1slice_model.Mm_per_ds
+ds_per_pixel = Mm_per_pixel / Mm_per_ds
+
+coords = np.stack(
+    np.meshgrid(np.linspace(x_min, x_max, np.round((x_max - x_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(y_min, y_max, np.round((y_max - y_min) / ds_per_pixel + 1).astype(int)),
+                np.linspace(0, 50 / Mm_per_ds, np.round(50 / Mm_per_pixel + 1).astype(int)),
+                indexing='ij'), -1)
+mfr_1slice_out = mfr_1slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+mfr_2slice_out = mfr_2slice_model.load_coords(coords, metrics=['j', 'b_nabla_bz'], progress=True)
+
+height_model = HeightTransformOutput('/glade/work/rjarolim/nf2/sst/13392_2slices_0851_v02/extrapolation_result.nf2')
+height_out = height_model.load_height_mapping()
+
+########################################################################################################################
+# y_slice = 470
+# y_range = [90, 130]
+
+y_slice = 150
+y_range = [50, 90]
+
+extent = np.array([x_min, x_max, y_min, y_max]) * Mm_per_ds
+
+fig, axs = plt.subplots(2, 1, figsize=(10, 5))
+
+ax = axs[0]
+im = ax.imshow(mfr_1slice_out['b'][:, :, 0, 2].to_value(u.G).T, origin='lower', cmap='gray',
+                 vmin=-2000, vmax=2000, extent=extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+
+ax = axs[1]
+im = ax.imshow(sst_data, origin='lower', cmap='gray', extent=extent)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Intensity')
+
+# draw red line at slice position
+[ax.plot([y_slice * Mm_per_pixel, y_slice * Mm_per_pixel], [y_range[0], y_range[1]], color='red', linestyle='--') for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+########################################################################################################################
+
+
+extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, 0, 70]
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+
+ax = axs[0]
+ax.imshow(mfr_1slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(mfr_2slice_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1, extent=extent)
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\hat{B} \cdot  \nabla \hat{B}_z$ [G Mm$^{-1}$]')
+
+[ax.set_xlabel('Y [Mm]') for ax in axs]
+axs[0].set_ylabel('Z [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+[ax.set_xlim(y_range) for ax in axs]
+[ax.set_ylim([0, 40]) for ax in axs]
+
+
+h_coords = height_out[0]['coords']
+h_y_slice = np.argmin(np.abs(h_coords[:, 0, 0, 0].to_value(u.Mm) - y_slice * Mm_per_pixel))
+ax.plot(np.linspace(y_min, y_max, h_coords.shape[1]) * Mm_per_ds, h_coords[h_y_slice, :, 0, 2].to_value(u.Mm), color='black', linestyle='--')
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b_nabla_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(mfr_1slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm, cmap='jet')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(np.linalg.norm(mfr_2slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm, cmap='jet')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|B|$ [G]')
+
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+[ax.set_xlim(y_range) for ax in axs]
+[ax.set_ylim([0, 40]) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b.png'), dpi=300, transparent=True)
+plt.close()
+
+
+
+##############################################################################
+
+extent = [x_min * Mm_per_ds, x_max * Mm_per_ds, y_min * Mm_per_ds, y_max * Mm_per_ds]
+j_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(mfr_1slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm, cmap='inferno')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(np.linalg.norm(mfr_2slice_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm, cmap='inferno')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G cm s$^{-1}$]')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_j.png'), dpi=300, transparent=True)
+plt.close()
+
+
+##############################################################################
+
+extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(mfr_1slice_out['metrics']['j'][y_slice, :, :].to_value(u.G / u.s), axis=-1).T, origin='lower', extent=extent, norm=b_norm, cmap='inferno')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(np.linalg.norm(mfr_2slice_out['metrics']['j'][y_slice, :, :].to_value(u.G / u.s), axis=-1).T, origin='lower', extent=extent, norm=b_norm, cmap='inferno')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G s$^{-1}$]')
+
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+[ax.set_xlim(y_range) for ax in axs]
+[ax.set_ylim([0, 40]) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_j_slice.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm(vmin=10)
+
+fig, ax = plt.subplots(1, 1, figsize=(5, 5))
+
+im = ax.imshow(np.linalg.norm(mfr_2slice_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower', extent=extent, norm=b_norm, cmap='jet')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|B|$ [G]')
+
+
+ax.set_xlabel('X [Mm]')
+ax.set_ylabel('Y [Mm]')
+ax.set_xlim([65 - 15, 65 + 15])
+ax.set_ylim([0, 20])
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b_multi.png'), dpi=300, transparent=True)
+plt.close()
\ No newline at end of file
diff --git a/nf2/evaluation/sst/compare_sst_eruption.py b/nf2/evaluation/sst/compare_sst_eruption.py
new file mode 100644
index 0000000..cc56649
--- /dev/null
+++ b/nf2/evaluation/sst/compare_sst_eruption.py
@@ -0,0 +1,208 @@
+import argparse
+import os
+
+import numpy as np
+from astropy import units as u
+from astropy.io import fits
+from matplotlib import pyplot as plt
+from matplotlib.colors import LogNorm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from nf2.evaluation.output import CartesianOutput
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Evaluate SST extrapolation.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+out_path = '/glade/work/rjarolim/nf2/sst/evaluation_eruption'
+os.makedirs(out_path, exist_ok=True)
+
+pre_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/sst/13392_2slices_0851_v01/extrapolation_result.nf2')
+post_model = CartesianOutput(
+    '/glade/work/rjarolim/nf2/sst/13392_2slices_1050_v03/extrapolation_result.nf2')
+
+Mm_per_pixel = 0.36
+pre_out = pre_model.load_cube(height_range=[0, 40], Mm_per_pixel=Mm_per_pixel, metrics=['j', 'b_nabla_bz'], progress=True)
+post_out = post_model.load_cube(height_range=[0, 40], Mm_per_pixel=Mm_per_pixel, metrics=['j', 'b_nabla_bz'], progress=True)
+
+########################################################################################################################
+y_slice_pre = 150
+y_slice_post = 200
+y_range = [50, 90]
+
+
+fig, axs = plt.subplots(2, 1, figsize=(10, 5))
+
+ax = axs[0]
+im = ax.imshow(pre_out['b'][:, :, 0, 2].to_value(u.G).T, origin='lower', cmap='gray',
+               vmin=-2000, vmax=2000)
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+ax.axvline(y_slice_pre, color='red', linestyle='--')
+
+ax = axs[1]
+im = ax.imshow(post_out['b'][:, :, 0, 2].to_value(u.G).T, origin='lower', cmap='gray')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='2%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label='Intensity')
+ax.axvline(y_slice_post, color='red', linestyle='--')
+
+# draw red line at slice position
+# [ax.plot([y_slice, y_slice], [y_range[0], y_range[1]], color='red', linestyle='--') for ax
+#  in axs]
+
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+########################################################################################################################
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(pre_out['metrics']['b_nabla_bz'][y_slice_pre, :, :].T, origin='lower',
+          cmap='RdBu_r', vmin=-.1, vmax=.1)
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(post_out['metrics']['b_nabla_bz'][y_slice_post, :, :].T, origin='lower',
+               cmap='RdBu_r', vmin=-.1, vmax=.1)
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\hat{B} \cdot  \nabla \hat{B}_z$ [G Mm$^{-1}$]')
+
+[ax.set_xlabel('Y [Mm]') for ax in axs]
+axs[0].set_ylabel('Z [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+# [ax.set_xlim(y_range) for ax in axs]
+# [ax.set_ylim([0, 40]) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b_nabla_bz.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+b_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(pre_out['b'][y_slice_pre, :, :].to_value(u.G), axis=-1).T, origin='lower',
+          norm=b_norm, cmap='jet')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(np.linalg.norm(post_out['b'][y_slice_post, :, :].to_value(u.G), axis=-1).T, origin='lower',
+               norm=b_norm, cmap='jet')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|B|$ [G]')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+# [ax.set_xlim(y_range) for ax in axs]
+# [ax.set_ylim([0, 40]) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+j_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(pre_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8,
+          origin='lower', norm=j_norm, cmap='inferno')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(
+    np.linalg.norm(post_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum(2).T * Mm_per_pixel * 1e8,
+    origin='lower', norm=j_norm, cmap='inferno')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G cm s$^{-1}$]')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_j.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+b_norm = LogNorm()
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+ax.imshow(np.linalg.norm(pre_out['metrics']['j'][y_slice_pre, :, :].to_value(u.G / u.s), axis=-1).T, origin='lower',
+          norm=b_norm, cmap='inferno')
+ax.set_title('NF2 - single height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+cax.set_axis_off()
+
+ax = axs[1]
+im = ax.imshow(np.linalg.norm(post_out['metrics']['j'][y_slice_post, :, :].to_value(u.G / u.s), axis=-1).T,
+               origin='lower', norm=b_norm, cmap='inferno')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|J|$ [G s$^{-1}$]')
+
+[ax.set_xlabel('X [Mm]') for ax in axs]
+axs[0].set_ylabel('Y [Mm]')
+[ax.axes.yaxis.set_ticklabels([]) for ax in axs[1:]]
+# [ax.set_xlim(y_range) for ax in axs]
+# [ax.set_ylim([0, 40]) for ax in axs]
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_j_slice.png'), dpi=300, transparent=True)
+plt.close()
+
+##############################################################################
+
+extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, 0, 70]
+b_norm = LogNorm(vmin=10)
+
+fig, ax = plt.subplots(1, 1, figsize=(5, 5))
+
+im = ax.imshow(np.linalg.norm(post_out['b'][y_slice, :, :].to_value(u.G), axis=-1).T, origin='lower',
+               extent=extent, norm=b_norm, cmap='jet')
+ax.set_title('NF2 - multi height')
+divider = make_axes_locatable(ax)
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|B|$ [G]')
+
+ax.set_xlabel('X [Mm]')
+ax.set_ylabel('Y [Mm]')
+ax.set_xlim([65 - 15, 65 + 15])
+ax.set_ylim([0, 20])
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_b_multi.png'), dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sst/compare_sst_full.py b/nf2/evaluation/sst/compare_sst_full.py
new file mode 100644
index 0000000..efea5c9
--- /dev/null
+++ b/nf2/evaluation/sst/compare_sst_full.py
@@ -0,0 +1,251 @@
+import argparse
+import os
+
+import numpy as np
+from astropy import units as u
+from astropy.io import fits
+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 energy, theta_J
+from nf2.evaluation.output import CartesianOutput, HeightTransformOutput
+
+# if __name__ == '__main__':
+parser = argparse.ArgumentParser(description='Evaluate SST extrapolation.')
+parser.add_argument('--output', type=str, help='output path.')
+args = parser.parse_args()
+
+out_path = '/glade/work/rjarolim/nf2/sst/evaluation'
+os.makedirs(out_path, exist_ok=True)
+
+# select intensity data in wings
+sst_data = fits.getdata('/glade/work/rjarolim/data/SST/panorama_8542_StkI.fits')[10]
+
+config1_file = '/glade/work/rjarolim/nf2/sst/13392_1slices_0851_v01/extrapolation_result.nf2'
+config2_file = '/glade/work/rjarolim/nf2/sst/13392_7699_0851_v01/extrapolation_result.nf2'
+config3_file = '/glade/work/rjarolim/nf2/sst/13392_2slices_0851_v02/extrapolation_result.nf2'
+config4_file = '/glade/work/rjarolim/nf2/sst/13392_3slices_0851_v04/extrapolation_result.nf2'
+config5_file = '/glade/work/rjarolim/nf2/sst/13392_4slices_0851_v01/extrapolation_result.nf2'
+
+config1_model = CartesianOutput(config1_file)
+
+config2_model = CartesianOutput(config2_file)
+
+config3_model = CartesianOutput(config3_file)
+
+config4_model = CartesianOutput(config4_file)
+
+config5_model = CartesianOutput(config5_file)
+
+
+config1_height_out = None
+
+config2_height_out = None
+
+config3_height_model = HeightTransformOutput(config3_file)
+config3_height_out = config3_height_model.load_height_mapping()
+
+config4_height_model = HeightTransformOutput(config4_file)
+config4_height_out = config4_height_model.load_height_mapping()
+
+config5_height_model = HeightTransformOutput(config5_file)
+config5_height_out = config5_height_model.load_height_mapping()
+
+x_min, x_max = config1_model.coord_range[0]
+y_min, y_max = config1_model.coord_range[1]
+
+Mm_per_pixel = 0.72
+height_range = [0, 40]
+Mm_per_ds = config1_model.Mm_per_ds
+ds_per_pixel = Mm_per_pixel / Mm_per_ds
+
+config1_out = config1_model.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range,
+                                      metrics=['j', 'b_nabla_bz', 'free_energy'], progress=True)
+config2_out = config2_model.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range,
+                                      metrics=['j', 'b_nabla_bz', 'free_energy'], progress=True)
+config3_out = config3_model.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range,
+                                      metrics=['j', 'b_nabla_bz', 'free_energy'], progress=True)
+config4_out = config4_model.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range,
+                                      metrics=['j', 'b_nabla_bz', 'free_energy'], progress=True)
+config5_out = config5_model.load_cube(Mm_per_pixel=Mm_per_pixel, height_range=height_range,
+                                      metrics=['j', 'b_nabla_bz', 'free_energy'], progress=True)
+
+
+
+config_outs = [config1_out, config2_out, config3_out, config4_out, config5_out]
+config_height_outs = [config1_height_out, config2_height_out, config3_height_out, config4_height_out, config5_height_out]
+labels = ['config1', 'config2', 'config3', 'config4', 'config5']
+# axs[0, 0].set_title('6173 Å')
+# axs[0, 1].set_title('+ 8542 Å')
+# axs[0, 2].set_title(r'+ 7699 Å')
+# axs[0, 3].set_title(r'+ 5896 Å ($B_\text{LOS}$)')
+########################################################################################################################
+y_slice = 150 // 2  # int(50 // Mm_per_pixel)
+y_range = [45, 85]
+j_norm = LogNorm(vmin=1e11, vmax=1e13)
+b_norm = LogNorm(vmin=10, vmax=1e3)
+extent = np.array([x_min, x_max, y_min, y_max]) * Mm_per_ds
+yz_extent = [y_min * Mm_per_ds, y_max * Mm_per_ds, *height_range]
+
+fig, axs = plt.subplots(4, 5, figsize=(20, 10))
+
+
+for i, (model_out, height_out) in enumerate(zip(config_outs, config_height_outs)):
+    column = axs[:, i]
+    #
+    ax = column[0]
+    im = ax.imshow(model_out['b'][:, :, 0, 2].to_value(u.G).T, origin='lower', cmap='gray',
+                   vmin=-2000, vmax=2000, extent=extent)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical', label='Bz [G]')
+    ax.plot([y_slice * Mm_per_pixel, y_slice * Mm_per_pixel], [y_range[0], y_range[1]], color='red', linestyle='--')
+    #
+    ax = column[1]
+    j_mag = np.linalg.norm(model_out['metrics']['j'].to_value(u.G / u.s), axis=-1)
+    im = ax.imshow(j_mag.sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=j_norm, cmap='inferno')
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\int{|J|}\,dz$ [G cm s$^{-1}$]')
+    #
+    ax = column[2]
+    im = ax.imshow(model_out['metrics']['b_nabla_bz'][y_slice, :, :].T, origin='lower',
+                   cmap='RdBu_r', vmin=-.1, vmax=.1, extent=yz_extent)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical', label=r'$\hat{B} \cdot  \nabla \hat{B}_z$ [G Mm$^{-1}$]')
+    ax.set_xlim(y_range)
+    ax.set_ylim([0, 30])
+    #
+    if height_out is not None:
+        for height_mapping in height_out:
+            h_coords = height_mapping['coords']
+            h_y_slice = np.argmin(np.abs(h_coords[:, 0, 0, 0].to_value(u.Mm) - y_slice * Mm_per_pixel))
+            ax.plot(np.linspace(y_min, y_max, h_coords.shape[1]) * Mm_per_ds, h_coords[h_y_slice, :, 0, 2].to_value(u.Mm),
+                    linestyle='--', color='black')
+    #
+    ax = column[3]
+    b_mag = np.linalg.norm(model_out['b'].to_value(u.G), axis=-1)
+    im = ax.imshow(b_mag[y_slice, :, :].T, origin='lower', extent=yz_extent, norm=b_norm, cmap='jet')
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical', label=r'$|B|$ [G]')
+    ax.set_xlim(y_range)
+    ax.set_ylim([0, 30])
+
+for i, label in enumerate(labels):
+    axs[0, i].set_title(label)
+
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_full_comparison.png'), dpi=300, transparent=True)
+plt.close()
+
+########################################################################################################################
+# plot energy and current density profiles
+
+fig, axs = plt.subplots(1, 4, figsize=(10, 10))
+
+
+heights = np.linspace(0, 40, config1_out['b'].shape[2])
+for i, model_out in enumerate(config_outs):
+    energy_profile = energy(model_out['b']).sum((0, 1)) * (Mm_per_pixel * 1e8) ** 2
+    free_energy_profile = model_out['metrics']['free_energy'].sum((0, 1)) * (Mm_per_pixel * 1e8) ** 2
+    current_profile = np.linalg.norm(model_out['metrics']['j'].to_value(u.G / u.s), axis=-1).sum((0, 1)) * (
+            Mm_per_pixel * 1e8) ** 2
+    b_mag = np.linalg.norm(model_out['b'].to_value(u.G), axis=-1)
+    jxb = np.linalg.norm(
+        np.cross(model_out['b'].to_value(u.G), model_out['metrics']['j'].to_value(u.G / u.s), axis=-1), axis=-1)
+    jxb_profile = (jxb / b_mag).sum((0, 1))
+    #
+    axs[0].plot(energy_profile, heights, label=labels[i])
+    axs[1].plot(free_energy_profile, heights, label=labels[i])
+    axs[2].plot(current_profile, heights, label=labels[i])
+    axs[3].plot(jxb_profile, heights, label=labels[i])
+
+ax = axs[0]
+ax.set_xlabel('Energy [erg cm$^{-1}$]')
+ax.set_ylabel('Height [Mm]')
+ax.set_title('Energy')
+ax.legend()
+ax.grid()
+ax.set_ylim([0, 40])
+
+ax = axs[1]
+ax.set_xlabel('Free energy [erg cm$^{-1}$]')
+ax.set_ylabel('Height [Mm]')
+ax.set_title('Free energy')
+ax.legend()
+ax.grid()
+ax.set_ylim([0, 40])
+
+ax = axs[2]
+ax.set_xlabel('Current density [G cm$^{2}$ s$^{-1}$]')
+ax.set_ylabel('Height [Mm]')
+ax.set_title('Current density')
+ax.legend()
+ax.grid()
+ax.set_ylim([0, 40])
+
+ax = axs[3]
+ax.set_xlabel('(J x B) / B [G cm$^{-1}$ s$^{-1}$]')
+ax.set_ylabel('Height [Mm]')
+ax.set_title('Force-freeness')
+ax.legend()
+ax.grid()
+ax.set_ylim([0, 40])
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_full_profiles.png'), dpi=300, transparent=True)
+plt.close()
+
+########################################################################################################################
+# print total energy stored in the field
+
+for config_out in config_outs:
+    energy_total = energy(config_out['b'].to_value(u.G)).sum() * (Mm_per_pixel * 1e8) ** 3
+    print(f'Energy: {energy_total * 1e-32:.2f} 10^32 erg')
+
+for config_out in config_outs:
+    free_energy_total = config_out['metrics']['free_energy'].sum() * (Mm_per_pixel * 1e8 * u.cm) ** 3
+    print(f'Free Energy: {free_energy_total.to_value(u.erg) * 1e-32:.2f} 10^32 erg')
+
+########################################################################################################################
+# print theta_J
+
+for config_out in config_outs:
+    theta = theta_J(config_out['b'].to_value(u.G), config_out['metrics']['j'].to_value(u.G / u.s))
+    print(f'Theta_J: {theta:.2f}')
+
+
+########################################################################################################################
+# plot height mapping
+
+extent = np.array([x_min, x_max, y_min, y_max]) * Mm_per_ds
+
+height_model = HeightTransformOutput(config5_file)
+height_out = height_model.load_height_mapping()
+
+fig, axs = plt.subplots(len(height_out) + 2, 1, figsize=(10, 15))
+
+j_mag = np.linalg.norm(config5_out['metrics']['j'].to_value(u.G / u.s), axis=-1)
+axs[0].imshow(sst_data, origin='lower', extent=extent, cmap='gray')
+axs[0].contour(j_mag.sum(2).T * Mm_per_pixel * 1e8, extent=extent, levels=[1e12], colors=['red'], linewidths=0.5)
+axs[0].grid()
+
+axs[1].imshow(j_mag.sum(2).T * Mm_per_pixel * 1e8, origin='lower', extent=extent, norm=LogNorm(1e11, 1e13),
+              cmap='plasma')
+
+for i, h_out in enumerate(height_out):
+    ax = axs[i + 2]
+    im = ax.imshow(h_out['coords'][..., 0, 2].to_value(u.Mm).T, origin='lower', extent=extent, cmap='jet',
+                   norm=LogNorm(1e-2, 10))
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='2%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical', label='Height [Mm]')
+    ax.grid()
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'height_mapping.png'), dpi=300, transparent=True)
+plt.close()
diff --git a/nf2/evaluation/sst/plot_data.py b/nf2/evaluation/sst/plot_data.py
new file mode 100644
index 0000000..3699ed3
--- /dev/null
+++ b/nf2/evaluation/sst/plot_data.py
@@ -0,0 +1,99 @@
+import glob
+import os
+
+from astropy.io import fits
+from matplotlib import pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+import numpy as np
+
+out_path = '/glade/work/rjarolim/nf2/sst/evaluation'
+os.makedirs(out_path, exist_ok=True)
+
+wl_keys = ['6173_StkIQUV_rebin', '7699_StkIQUV_rebin_v2', '8542_StkIQUV_rebin']
+
+los_files = [f'/glade/work/rjarolim/data/SST/campaign_2023_v2/BLOS_panorama0851_{i}.fits' for i in wl_keys]
+trv_files = [f'/glade/work/rjarolim/data/SST/campaign_2023_v2/BTRV_panorama0851_{i}.fits' for i in wl_keys]
+azi_files = [f'/glade/work/rjarolim/data/SST/campaign_2023_v2/BAZI_panorama0851_{i}.fits' for i in wl_keys]
+mask = fits.getdata('/glade/work/rjarolim/data/SST/campaign_2023_v2/mask.fits')
+
+fig, axs = plt.subplots(len(wl_keys), 4, figsize=(13, 6))
+
+labels = ['6173 Å', '7699 Å' , '8542 Å']
+
+# Plot LOS
+for i in range(len(labels)):
+    ax = axs[i, 0]
+    data_los = fits.getdata(los_files[i])
+    data_los[mask == 1] = np.nan
+    im = ax.imshow(data_los, cmap='gray', origin='lower', 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')
+    ax.set_title(f'LOS {labels[i]}')
+    # Plot TRV
+    ax = axs[i, 1]
+    data_trv = fits.getdata(trv_files[i])
+    data_trv[mask == 1] = np.nan
+    im = ax.imshow(data_trv, cmap='gray', origin='lower', 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')
+    ax.set_title(f'TRV {labels[i]}')
+    # Plot AZI
+    ax = axs[i, 2]
+    data_azi = fits.getdata(azi_files[i])
+    data_azi[mask == 1] = np.nan
+    im = ax.imshow(data_azi % np.pi, cmap='twilight', origin='lower', vmin=0, vmax=np.pi)
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes('right', size='5%', pad=0.05)
+    fig.colorbar(im, cax=cax, orientation='vertical')
+    ax.set_title(f'AZI {labels[i]}')
+    #
+    ax = axs[i, 3]
+    b_abs = np.sqrt(data_los**2 + data_trv**2)
+    ax.hist(b_abs.flatten(), bins=100, range=(0, 1000))
+    ax.set_title(r'$\left\Vert B \right\Vert$ histogram ' + f'{labels[i]}')
+    # add mean
+    mean_b = np.nanmean(b_abs)
+    mean_b_los = np.nanmean(np.abs(data_los))
+    ax.axvline(mean_b, color='red', linestyle='--', label=r'$\overline{\left\Vert B \right\Vert}$: ' + f'{mean_b:.0f} G')
+    ax.axvline(mean_b_los, color='blue', linestyle='--', label=r'$\overline{\left\Vert B_\text{LOS} \right\Vert}$: ' + f'{mean_b_los:.0f} G')
+    ax.axvline(200, color='black', linestyle='--', label=r'$\left\Vert B \right\Vert = 200$ G', alpha=0.5)
+    ax.legend()
+    ax.set_ylim(0, 1.5e5)
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_data_overview.png'))
+plt.show()
+
+########################################################################################################################
+# scatter plots
+
+b0_mag = np.sqrt(fits.getdata(los_files[0])**2 + fits.getdata(trv_files[0])**2)
+b1_mag = np.sqrt(fits.getdata(los_files[1])**2 + fits.getdata(trv_files[1])**2)
+b2_mag = np.sqrt(fits.getdata(los_files[2])**2 + fits.getdata(trv_files[2])**2)
+
+fig, axs = plt.subplots(1, 2, figsize=(10, 5))
+
+ax = axs[0]
+# 2d histogram
+ax.hist2d(b0_mag.flatten(), b1_mag.flatten(), bins=100, range=[[0, 3000], [0, 3000]], cmap='viridis', norm='log')
+ax.set_xlabel(r'$\left\Vert B \right\Vert$ 6173 Å [G]')
+ax.set_ylabel(r'$\left\Vert B \right\Vert$ 7699 Å [G]')
+
+
+ax = axs[1]
+ax.hist2d(b0_mag.flatten(), b2_mag.flatten(), bins=100, range=[[0, 3000], [0, 3000]], cmap='viridis', norm='log')
+ax.set_xlabel(r'$\left\Vert B \right\Vert$ 6173 Å [G]')
+ax.set_ylabel(r'$\left\Vert B \right\Vert$ 8542 Å [G]')
+
+for ax in axs:
+    ax.plot([0, 3000], [0, 3000], color='black', linestyle='--')
+    ax.set_xlim(0, 3000)
+    ax.set_ylim(0, 3000)
+    ax.set_aspect('equal')
+
+plt.tight_layout()
+plt.savefig(os.path.join(out_path, 'sst_scatter.png'))
+plt.close()
+
diff --git a/nf2/extrapolate.py b/nf2/extrapolate.py
index 81db90c..116c1e2 100644
--- a/nf2/extrapolate.py
+++ b/nf2/extrapolate.py
@@ -3,7 +3,6 @@
 import shutil
 
 import torch
-import yaml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint, LambdaCallback
 from pytorch_lightning.loggers import WandbLogger
@@ -16,9 +15,10 @@
 from nf2.loader.vsm import VSMDataModule
 from nf2.train.mapping import load_callbacks
 from nf2.train.module import NF2Module, save
+from nf2.train.util import load_yaml_config
 
 
-def run(base_path, data, work_directory=None, logging={}, model={}, training={}, config=None):
+def run(base_path, data, work_directory=None, callbacks=[], 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.
@@ -39,14 +39,14 @@ def run(base_path, data, work_directory=None, logging={}, model={}, training={},
     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
+    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}
+                   'model': model, 'training': training, 'config': config, 'data': data}
     wandb_logger.experiment.config.update(config_dict, allow_val_change=True)
 
     # restore model checkpoint from wandb
@@ -73,7 +73,7 @@ def run(base_path, data, work_directory=None, logging={}, model={}, training={},
         raise NotImplementedError(f'Unknown data loader {data["type"]}')
 
     # initialize callbacks
-    callbacks = load_callbacks(data_module)
+    callback_modules = load_callbacks(data_module, additional_callbacks=callbacks)
 
     nf2 = NF2Module(data_module.validation_dataset_mapping, data_module.config,
                     model_kwargs=model, **training)
@@ -92,7 +92,7 @@ def run(base_path, data, work_directory=None, logging={}, model={}, training={},
 
     torch.set_float32_matmul_precision('medium')  # for A100 GPUs
     n_gpus = torch.cuda.device_count()
-    callbacks += [checkpoint_callback, save_callback]
+    callback_modules += [checkpoint_callback, save_callback]
 
     trainer = Trainer(max_epochs=max_epochs,
                       logger=wandb_logger,
@@ -103,7 +103,7 @@ def run(base_path, data, work_directory=None, logging={}, model={}, training={},
                       val_check_interval=val_check_interval,
                       check_val_every_n_epoch=val_every_n_epochs,
                       gradient_clip_val=0.1,
-                      callbacks=callbacks)
+                      callbacks=callback_modules)
 
     trainer.fit(nf2, data_module, ckpt_path='last')
     save(save_path, nf2, data_module, config_dict)
@@ -115,14 +115,12 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--config', type=str, required=True,
                         help='config file for the simulation')
-    args = parser.parse_args()
+    args, overwrite_args = parser.parse_known_args()
 
-    with open(args.config) as config:
-        info = yaml.safe_load(config)
-        for key, value in info.items():
-            args.__dict__[key] = value
+    yaml_config_file = args.config
+    config = load_yaml_config(yaml_config_file, overwrite_args)
 
-    run(**args.__dict__)
+    run(**config)
 
 
 if __name__ == '__main__':
diff --git a/nf2/extrapolate_series.py b/nf2/extrapolate_series.py
index 5dc8f6f..2260c38 100644
--- a/nf2/extrapolate_series.py
+++ b/nf2/extrapolate_series.py
@@ -4,7 +4,6 @@
 import shutil
 
 import torch
-import yaml
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import LambdaCallback
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -14,6 +13,7 @@
 from nf2.loader.spherical import SphericalSeriesDataModule
 from nf2.train.mapping import load_callbacks
 from nf2.train.module import NF2Module, save
+from nf2.train.util import load_yaml_config
 
 
 def run(base_path, data, meta_path, work_directory=None, logging={}, model={}, training={}, config=None):
@@ -55,7 +55,7 @@ def run(base_path, data, meta_path, work_directory=None, logging={}, model={}, t
         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)
+    fits_paths, error_paths = _load_paths(data['data_path'])
 
     # reload model training
     ckpts = sorted(glob.glob(os.path.join(base_path, '*.nf2')))
@@ -105,9 +105,12 @@ def run(base_path, data, meta_path, work_directory=None, logging={}, model={}, t
     trainer.fit(nf2, data_module, ckpt_path=ckpt_path)
 
 
-def _load_paths(data):
-    data_path = data['data_path']
-    if isinstance(data_path, str):
+def _load_paths(data_path):
+    if isinstance(data_path, list):
+        results = [_load_paths(d) for d in data_path]
+        fits_paths = [f for r in results for f in r[0]]
+        error_paths = [f for r in results for f in r[1]] if all([r[1] is not None for r in results]) else None
+    elif 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
@@ -158,15 +161,13 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('--config', type=str, required=True,
                         help='config file for the simulation')
-    args = parser.parse_args()
+    args, overwrite_args = parser.parse_known_args()
 
-    with open(args.config) as config:
-        info = yaml.safe_load(config)
-        for key, value in info.items():
-            args.__dict__[key] = value
+    yaml_config_file = args.config
+    config = load_yaml_config(yaml_config_file, overwrite_args)
 
-    run(**args.__dict__)
+    run(**config)
 
 
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()
diff --git a/nf2/loader/base.py b/nf2/loader/base.py
index 8839d1e..f90afb5 100644
--- a/nf2/loader/base.py
+++ b/nf2/loader/base.py
@@ -93,12 +93,14 @@ 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,
+                 bin=1, height_mapping=None, plot=True, los_trv_azi=False, ambiguous_azimuth=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)
+        self.bz = np.copy(b[..., 0]) if los_trv_azi else np.copy(b[..., 2])
+
         # normalize
         if los_trv_azi:
             b[..., :2] /= G_per_dB
@@ -152,6 +154,9 @@ def __init__(self, b, b_err=None, coords=None,
             if plot and not los_trv_azi:
                 _plot_B(b * G_per_dB, coords * Mm_per_ds)
 
+        # prepare azimuth data after plotting
+        if los_trv_azi and ambiguous_azimuth:
+            b[..., 2] = np.mod(b[..., 2], np.pi)
 
         tensors = {k: v.reshape((-1, *v.shape[2:])).astype(np.float32) for k, v in tensors.items()}
 
diff --git a/nf2/loader/fits.py b/nf2/loader/fits.py
index 650b00a..88fbd53 100644
--- a/nf2/loader/fits.py
+++ b/nf2/loader/fits.py
@@ -17,7 +17,7 @@ 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,
+                 Mm_per_ds=.36 * 320, G_per_dB=2500, z_range=None, 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]
@@ -33,8 +33,9 @@ def __init__(self, slices, work_directory, boundary_config=None,
 
         # 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)
+        z_range = [0, 100] if z_range is None else z_range
+        z_range_arr = np.array([z_range]) / Mm_per_ds
+        coord_range = np.concatenate([coord_range, z_range_arr], axis=0)
         random_config = random_config if random_config is not None else {}
         random_dataset = RandomCoordinateDataset(coord_range, **random_config)
 
@@ -62,15 +63,12 @@ def __init__(self, slices, work_directory, boundary_config=None,
 
         # 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}
+            {'type': 'potential', 'strides': 4, 'Mm_per_pixel': ds_per_pixel * Mm_per_ds}
         if boundary_config['type'] == 'potential':
-            bz = fits.getdata(boundary_config['file']).transpose()
+            bz = bottom_boundary_dataset.bz
             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'],
+                                                         height_pixel=coord_range[2, -1] / ds_per_pixel,
                                                          ds_per_pixel=ds_per_pixel, G_per_dB=G_per_dB,
                                                          work_directory=work_directory,
                                                          strides=boundary_config['strides'])
@@ -80,7 +78,8 @@ def __init__(self, slices, work_directory, boundary_config=None,
         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,
+            valid_slice_config = {k: v for k, v in slice_config.items() if k != 'batch_size'}
+            validation_boundary_dataset = self.init_boundary_dataset(**valid_slice_config,
                                                                      Mm_per_ds=Mm_per_ds, G_per_dB=G_per_dB,
                                                                      shuffle=False, filter_nans=False,
                                                                      work_directory=work_directory,
@@ -93,8 +92,8 @@ def __init__(self, slices, work_directory, boundary_config=None,
         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,
+        config = {'type': 'cartesian', 'max_height': z_range[1],
+                  'Mm_per_ds': Mm_per_ds, 'G_per_dB': G_per_dB,
                   'coord_range': [], 'ds_per_pixel': [], 'height_mapping': [], 'wcs': []}
         for ds in slice_datasets:
             config['coord_range'].append(ds.coord_range)
@@ -111,6 +110,8 @@ def init_boundary_dataset(self, type='fits', **kwargs):
             return LosTrvAziFITSDataset(**kwargs)
         elif type == 'sharp':
             return SHARPDataset(**kwargs)
+        elif type == 'fld_inc_azi':
+            return FldIncAziFITSDataset(**kwargs)
         else:
             raise ValueError(f'Unknown boundary type: {type}')
 
@@ -162,7 +163,7 @@ def __init__(self, fits_path, error_path=None, coords_path=None,
 
 class LosTrvAziFITSDataset(MapDataset):
 
-    def __init__(self, fits_path,
+    def __init__(self, fits_path, mask_path=None,
                  bin=1, slice=None, load_map=True, **kwargs):
         file_los = fits_path['B_los']
         file_trv = fits_path['B_trv']
@@ -177,10 +178,23 @@ def __init__(self, fits_path,
             trv_data = trv_map.data
             azi_data = azi_map.data
             wcs = los_map.wcs
+            if mask_path is not None:
+                mask_map = Map(mask_path)
+                mask_map = process_map(mask_map, slice, bin)
+                mask = mask_map.data
+                los_data[mask] = 0 # np.nan
+                trv_data[mask] = 0 # np.nan
+                azi_data[mask] = 0 # np.nan
         else:
             los_data = fits.getdata(file_los)
             trv_data = fits.getdata(file_trv)
             azi_data = fits.getdata(file_azi)
+            if mask_path is not None:
+                mask = fits.getdata(mask_path)
+                mask = np.array(mask, dtype=bool)
+                los_data[mask] = 0 # np.nan
+                trv_data[mask] = 0 # np.nan
+                azi_data[mask] = 0 # np.nan
             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]]
@@ -195,6 +209,47 @@ def __init__(self, fits_path,
 
         super().__init__(b=b, wcs=wcs, los_trv_azi=True, **kwargs)
 
+class FldIncAziFITSDataset(MapDataset):
+
+    def __init__(self, fits_path, bin=1, slice=None, **kwargs):
+        file_fld = fits_path['B_fld']
+        file_inc = fits_path['B_inc']
+        file_azi = fits_path['B_azi']
+
+        fld_map, inc_map, azi_map = Map(file_fld), Map(file_inc), Map(file_azi)
+        fld_map.data[:] = np.flip(fld_map.data, axis=(0, 1))
+        inc_map.data[:] = np.flip(inc_map.data, axis=(0, 1))
+        azi_map.data[:] = np.flip(azi_map.data, axis=(0, 1))
+
+        fld_map = process_map(fld_map, slice, bin)
+        inc_map = process_map(inc_map, slice, bin)
+        azi_map = process_map(azi_map, slice, bin)
+
+        fld = fld_map.data
+        inc = np.deg2rad(inc_map.data)
+        azi = np.deg2rad(azi_map.data)
+
+        # apply disambiguation
+        if 'B_amb' in fits_path:
+            file_amb = fits_path['B_amb']
+            amb_map = Map(file_amb)
+            amb_map.data[:] = np.flip(amb_map.data, axis=(0, 1))
+            amb_map = process_map(amb_map, slice, bin)
+            amb = amb_map.data
+
+            amb_weak = 2
+            condition = (amb.astype(int) >> amb_weak).astype(bool)
+            azi[condition] += np.pi
+
+        wcs = fld_map.wcs
+
+        los = fld * np.cos(inc)
+        trv = fld * np.sin(inc)
+
+        b = np.stack([los, trv, (np.pi - azi) % (2 * np.pi)]).transpose()
+
+        super().__init__(b=b, wcs=wcs, los_trv_azi=True, **kwargs)
+
 
 class FITSSeriesDataModule(FITSDataModule):
 
@@ -210,19 +265,24 @@ def __init__(self, fits_paths, error_paths=None, *args, **kwargs):
         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)
+        # skip reload if already initialized - for initial epoch
+        if self.initialized:
+            self.initialized = False
+            print('Currently loaded:', self.current_id)
+            return super().train_dataloader()
+        # 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
+        print('Currently loaded:', self.current_id)
         return super().train_dataloader()
 
 
+
 class SHARPDataset(FITSDataset):
 
     def __init__(self, **kwargs):
@@ -231,7 +291,7 @@ def __init__(self, **kwargs):
 
 class PotentialBoundaryDataset(TensorsDataset):
 
-    def __init__(self, bz, height_pixel, ds_per_pixel, G_per_dB, strides=2, batch_size=2 ** 10, **kwargs):
+    def __init__(self, bz, height_pixel, ds_per_pixel, G_per_dB, strides=2, batch_size=2 ** 12, **kwargs):
         coords, b_err, b = load_potential_field_data(bz, height_pixel, strides,
                                                      progress=False)
         coords = coords * ds_per_pixel
diff --git a/nf2/loader/muram.py b/nf2/loader/muram.py
index 37195e0..3b91584 100644
--- a/nf2/loader/muram.py
+++ b/nf2/loader/muram.py
@@ -16,12 +16,19 @@ def __init__(self, slices, work_directory, boundary_config=None, random_config=N
                  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:
+        slice_datasets = {}
+        bottom_boundary_dataset = None
+        for i, slice_config in enumerate(slices):
+            ds_id = slice_config.pop('name', f'boundary_{i + 1:02d}')
+            bottom_boundary = slice_config.pop('bottom', False)
             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)
+            slice_datasets[ds_id] = muram_dataset
+            if bottom_boundary:
+                bottom_boundary_dataset = muram_dataset
 
-        bottom_boundary_dataset = slice_datasets[0]
+        #
+        if bottom_boundary_dataset is None:
+            bottom_boundary_dataset = list(slice_datasets.values())[0]
 
         # random sampling dataset
         coord_range = bottom_boundary_dataset.coord_range
@@ -48,15 +55,14 @@ def __init__(self, slices, work_directory, boundary_config=None, random_config=N
             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
+        for ds_id, dataset in slice_datasets.items():
+            training_datasets[ds_id] = 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}
+        boundary_config = boundary_config if boundary_config is not None else {'type': 'potential', 'strides': 4}
         if boundary_config['type'] == 'potential':
-            sl, Nvar, shape, time = read_muram_slice(slices[0]['data_path'])
-            bz = sl[5, :, :] * np.sqrt(4 * np.pi)
+            bz = bottom_boundary_dataset.bz
             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,
@@ -82,7 +88,7 @@ def __init__(self, slices, work_directory, boundary_config=None, random_config=N
         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:
+        for ds in slice_datasets.values():
             config['coord_range'].append(ds.coord_range)
             config['ds_per_pixel'].append(ds.ds_per_pixel)
             config['height_mapping'].append(ds.height_mapping)
@@ -101,7 +107,7 @@ def __init__(self, slices, work_directory, boundary_config=None, random_config=N
 
 class MURaMDataset(MapDataset):
 
-    def __init__(self, data_path, *args, **kwargs):
+    def __init__(self, data_path, los_trv_azi_transform=False, scaling=None, *args, **kwargs):
         sl, Nvar, shape, time = read_muram_slice(data_path)
 
         bz = sl[5, :, :] * np.sqrt(4 * np.pi)
@@ -110,21 +116,13 @@ def __init__(self, data_path, *args, **kwargs):
 
         b = np.stack([bx, by, bz], axis=-1)
 
-        super().__init__(b, Mm_per_pixel=0.192, *args, **kwargs)
-
-class MURaMLosTrvAziDataset(MapDataset):
+        if scaling is not None:
+            b = b * scaling
 
-    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)
+        if los_trv_azi_transform:
+            b = img_to_los_trv_azi(b, f=np)
 
-        super().__init__(b, Mm_per_pixel=0.192, los_trv_azi=True, *args, **kwargs)
+        super().__init__(b, Mm_per_pixel=0.192, los_trv_azi=los_trv_azi_transform, *args, **kwargs)
 
 class MURaMSnapshot():
 
@@ -158,7 +156,7 @@ def B(self):
     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):
+    def load_cube(self, resolution=0.192 * u.Mm / u.pix, height=100 * u.Mm, target_tau=1):
 
         b = self.B
         tau = self.tau
@@ -174,11 +172,11 @@ def load_cube(self, resolution=0.192 * u.Mm / u.pix, height=100 * u.Mm):
         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
+        pix_height = np.argmin(np.abs(tau - target_tau), 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))
+        max_height = min_height + int((height / resolution).to_value(u.pix))
         b = b[:, :, min_height:max_height]
         tau = tau[:, :, min_height:max_height]
 
diff --git a/nf2/loader/spherical.py b/nf2/loader/spherical.py
index d42d96e..f33d01d 100644
--- a/nf2/loader/spherical.py
+++ b/nf2/loader/spherical.py
@@ -1,6 +1,6 @@
 import os
 from collections import OrderedDict
-from copy import copy
+from copy import copy, deepcopy
 
 import numpy as np
 import pfsspy
@@ -31,12 +31,12 @@ def __init__(self, b, coords, spherical_coords, G_per_dB, Mm_per_ds,
             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)
+        b = b.reshape((-1, 3))
+        coords = coords.reshape((-1, 3))
         if b_err is not None:
-            b_err = np.concatenate([e.reshape((-1, 3)) for e in b_err]).astype(np.float32)
+            b_err = b_err.reshape((-1, 3))
         if transform is not None:
-            transform = np.concatenate([t.reshape((-1, 3, 3)) for t in transform]).astype(np.float32)
+            transform = transform.reshape((-1, 3, 3))
 
         # normalize data
         b /= G_per_dB
@@ -93,7 +93,7 @@ def _plot(self, b, coords, spherical_coords):
 
 class SphericalMapDataset(SphericalSliceDataset):
 
-    def __init__(self, files, mask_configs=[], **kwargs):
+    def __init__(self, files, mask_configs=None, insert=None, **kwargs):
         # load maps
         r_map = Map(files['Br'])
         t_map = Map(files['Bt'])
@@ -126,7 +126,32 @@ def __init__(self, files, mask_configs=[], **kwargs):
         else:
             b_error_spherical = None
 
+        # 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 insert_config in insert:
+            insert_br = Map(insert_config['Br'])
+            insert_bt = Map(insert_config['Bt'])
+            insert_bp = Map(insert_config['Bp'])
+
+            # reproject to reference map
+            ref_map = Map(r_map.data, r_map.meta)
+            # use reference time to avoid temporal shift
+            ref_map.meta['date-obs'] = insert_br.date.to_datetime().isoformat()
+            ref_wcs = ref_map.wcs
+            insert_br = insert_br.reproject_to(ref_wcs)
+            insert_bt = insert_bt.reproject_to(ref_wcs)
+            insert_bp = insert_bp.reproject_to(ref_wcs)
+
+            insert_b_spherical = np.stack([insert_br.data, insert_bt.data, insert_bp.data]).transpose()
+            insert_b_cartesian = vector_spherical_to_cartesian(insert_b_spherical, spherical_coords)
+
+            nan_mask = ~np.isnan(insert_b_spherical)
+            b_spherical[nan_mask] = insert_b_spherical[nan_mask]
+            b_cartesian[nan_mask] = insert_b_cartesian[nan_mask]
+
         # apply masking
+        mask_configs = mask_configs if mask_configs is not None else []
         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,
@@ -143,6 +168,7 @@ def __init__(self, files, mask_configs=[], **kwargs):
             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]
+            b_cartesian = b_cartesian[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:
@@ -221,6 +247,8 @@ def __init__(self, Br, radius_range, source_surface_height=2.5, resample=[360, 1
         insert = insert if isinstance(insert, list) else [insert]
         for file_in in insert:
             map_in = Map(file_in)
+            # prevent temporal rotation
+            potential_r_map.meta['date-obs'] = map_in.date.to_datetime().isoformat()
             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]
@@ -258,13 +286,16 @@ def __init__(self, Br, radius_range, source_surface_height=2.5, resample=[360, 1
         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,
+                 max_radius=1.3,
+                 Mm_per_ds= (1 * u.solRad).to_value(u.Mm),
+                 G_per_dB=None, work_directory=None,
                  batch_size=4096, **kwargs):
 
         self.ds_mapping = {'map': SphericalMapDataset,
@@ -274,7 +305,6 @@ def __init__(self, train_configs, validation_configs,
                            '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]
@@ -297,6 +327,7 @@ def __init__(self, train_configs, validation_configs,
     def load_config(self, configs, general_config, prefix='train'):
         datasets = OrderedDict()
         for i, config in enumerate(configs):
+            config = deepcopy(config)
             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
@@ -311,20 +342,55 @@ def load_config(self, configs, general_config, prefix='train'):
 
 
 class SphericalSeriesDataModule(SphericalDataModule):
-    def __init__(self, full_disk_files, *args, **kwargs):
+
+    def __init__(self, fits_paths, synoptic_fits_path, *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)
+        self.fits_paths = copy(fits_paths)
+        self.synoptic_fits_path = synoptic_fits_path
+
+        self.current_id = os.path.basename(self.fits_paths[0]['Br']).split('.')[-3]
+
+        self.initialized = True  # only required for first iteration
+        train_configs = self._build_config(fits_paths[0])
+        super().__init__(train_configs, *self.args, **self.kwargs)
+
+    def _build_config(self, fits):
+        full_disk_config = {
+            'type': 'map',
+            'ds_id': 'full_disk',
+            'batch_size': 4096,
+            'files': {'Br': fits['Br'], 'Bt': fits['Bt'], 'Bp': fits['Bp'], }
+        }
+        synoptic_config = {
+            'type': 'map',
+            'ds_id': 'synoptic',
+            'batch_size': 4096,
+            'files': {'Br': self.synoptic_fits_path['Br'],
+                      'Bt': self.synoptic_fits_path['Bt'],
+                      'Bp': self.synoptic_fits_path['Bp']},
+            'mask_configs': [{'type': 'reference', 'file': fits['Br']}]
+        }
+        random_config = {
+            'type': 'random_spherical',
+            'ds_id': 'random',
+            'batch_size': 16384
+        }
+        train_configs = [full_disk_config, synoptic_config, random_config]
+        return train_configs
 
     def train_dataloader(self):
-        if len(self.full_disk_files) == 0:
-            return None
+        # skip reload if already initialized - for initial epoch
+        if self.initialized:
+            self.initialized = False
+            print('Currently loaded:', self.current_id)
+            return super().train_dataloader()
+        # update ID
+        self.current_id = os.path.basename(self.fits_paths[0]['Br']).split('.')[-3]
         # 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]
+        train_configs = self._build_config(self.fits_paths[0])
+        super().__init__(train_configs, *self.args, **self.kwargs)
+        # continue with next file in list
+        del self.fits_paths[0]
+        print('Currently loaded:', self.current_id)
         return super().train_dataloader()
diff --git a/nf2/train/callback.py b/nf2/train/callback.py
index 28518e1..808a959 100644
--- a/nf2/train/callback.py
+++ b/nf2/train/callback.py
@@ -373,6 +373,77 @@ def plot_b_coords(self, b, b_true, original_coords, transformed_coords):
         wandb.log({f"{self.validation_dataset_key} - B": fig})
         plt.close('all')
 
+class DisambiguationCallback(Callback):
+
+    def __init__(self, validation_dataset_key, cube_shape, gauss_per_dB, Mm_per_ds, name=None):
+        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
+        self.name = name
+
+    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_true = outputs['b_true']
+        flip = outputs['flip']
+
+        flip = flip.cpu().numpy().reshape([*self.cube_shape])
+        azimuth_amb = b_true[..., 2].cpu().numpy().reshape([*self.cube_shape])
+
+        azimuth = azimuth_amb + np.round(flip) * np.pi
+
+        extent = None
+
+        fig, axs = plt.subplots(3, 2, figsize=(8, 8))
+
+        ax = axs[0, 0]
+        im = ax.imshow(azimuth.T, cmap='twilight', origin='lower',
+                         extent=extent, vmin=0, vmax=2 * np.pi)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Azimuth [rad]')
+
+        ax = axs[0, 1]
+        im = ax.imshow(azimuth_amb.T, cmap='twilight', origin='lower',
+                         extent=extent, vmin=0, vmax=2 * np.pi)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Azimuth [rad]')
+
+        ax = axs[1, 0]
+        im = ax.imshow(azimuth.T % (np.pi), cmap='twilight', origin='lower',
+                         extent=extent, vmin=0, vmax=np.pi)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Azimuth [rad]')
+
+        ax = axs[1, 1]
+        im = ax.imshow(azimuth_amb.T % (np.pi), cmap='twilight', origin='lower',
+                            extent=extent, vmin=0, vmax=np.pi)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Azimuth [rad]')
+
+        ax = axs[2, 0]
+        im = ax.imshow(flip.T, cmap='bwr', origin='lower', extent=extent, vmin=0, vmax=1)
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        plt.colorbar(im, cax=cax, label='Flip Probability')
+
+        axs[2, 1].set_axis_off()
+
+        axs[0, 0].set_title('Predicted Azimuth')
+        axs[0, 1].set_title('Ambiguous Azimuth')
+
+        fig.tight_layout()
+        name = f"{self.validation_dataset_key} - Disambiguation" if self.name is None else self.name
+        wandb.log({name: fig})
+        plt.close('all')
+
+
 class LosTrvAziBoundaryCallback(Callback):
 
     def __init__(self, validation_dataset_key, cube_shape, gauss_per_dB, Mm_per_ds):
@@ -386,44 +457,89 @@ def on_validation_epoch_end(self, trainer, pl_module):
             return
 
         outputs = pl_module.validation_outputs[self.validation_dataset_key]
-        b = outputs['b']
+        b_xyz = 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_xyz = torch.einsum('ijk,ik->ij', transform, b_xyz)
 
-        b = img_to_los_trv_azi(b, f=torch)
-        b[..., 0:2] = b[..., 0:2] * self.gauss_per_dB
+        b_xyz = b_xyz * self.gauss_per_dB
+        b_los_trv_azi = img_to_los_trv_azi(b_xyz, f=torch)
         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 = 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)
+        b_xyz_amb = los_trv_azi_to_img(b_los_trv_azi, ambiguous=True, f=torch)
+        b_true_xyz_amb = 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.abs(b_xyz_amb - b_true_xyz_amb)
         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_los_trv_azi = b_los_trv_azi.cpu().numpy().reshape([*self.cube_shape, 3])
         b_true = b_true.cpu().numpy().reshape([*self.cube_shape, 3])
+        b_xyz = b_xyz.cpu().numpy().reshape([*self.cube_shape, 3])
+        b_true_xyz = b_true_xyz.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)
+            self.plot_b_coords(b_los_trv_azi, 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)
+            self.plot_b(b_los_trv_azi, b_true, original_coords)
+            #
+            self.plot_bxyz(b_xyz, b_true_xyz, original_coords)
+
+    def plot_bxyz(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} - Bxyz": fig})
+        plt.close('all')
 
     def plot_b(self, b, b_true, original_coords):
         extent = [original_coords[..., 0].min(), original_coords[..., 0].max(),
diff --git a/nf2/train/loss.py b/nf2/train/loss.py
index 5151c2a..8f8bf39 100644
--- a/nf2/train/loss.py
+++ b/nf2/train/loss.py
@@ -1,8 +1,10 @@
 import torch
+from astropy import units as u
 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
+from nf2.train.model import jacobian
+
 
 class BaseLoss(nn.Module):
 
@@ -31,8 +33,7 @@ def forward(self, b, jac_matrix, *args, **kwargs):
         #
         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()
+        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7)
         force_free_loss = torch.sum(jxb ** 2, dim=-1) / normalization
         #
         return force_free_loss.mean()
@@ -64,11 +65,62 @@ def forward(self, b, jac_matrix, *args, **kwargs):
         #
         equation = jxb - grad_P
         # assure that the normalization does not influence the loss
-        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7).detach()
+        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7)
         loss = torch.sum(equation ** 2, dim=-1) / normalization
         #
         return loss.mean()
 
+class ImplicitMagnetoStaticLoss(BaseLoss):
+
+    def __init__(self, divergence_regularization=True, **kwargs):
+        super().__init__(**kwargs)
+        self.divergence_regularization = divergence_regularization
+
+    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]
+        #
+        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)
+        jac_matrix = jacobian(jxb, coords)
+        dJxBx_dx = jac_matrix[:, 0, 0]
+        dJxBy_dx = jac_matrix[:, 1, 0]
+        dJxBz_dx = jac_matrix[:, 2, 0]
+        dJxBx_dy = jac_matrix[:, 0, 1]
+        dJxBy_dy = jac_matrix[:, 1, 1]
+        dJxBz_dy = jac_matrix[:, 2, 1]
+        dJxBx_dz = jac_matrix[:, 0, 2]
+        dJxBy_dz = jac_matrix[:, 1, 2]
+        dJxBz_dz = jac_matrix[:, 2, 2]
+        #
+        rot_x = dJxBz_dy - dJxBy_dz
+        rot_y = dJxBx_dz - dJxBz_dx
+        rot_z = dJxBy_dx - dJxBx_dy
+        #
+        curl_JxB = torch.stack([rot_x, rot_y, rot_z], -1)
+        div_JxB = dJxBx_dx + dJxBy_dy + dJxBz_dz
+        # assure that the normalization does not influence the loss
+        normalization = (torch.sum(b ** 2, dim=-1) + 1e-7)
+        curl_JxB_loss = torch.sum(curl_JxB ** 2, dim=-1)
+        div_JxB_loss = div_JxB.pow(2)
+        if self.divergence_regularization:
+            loss = curl_JxB_loss + div_JxB_loss * (torch.exp(coords[:, 2] * 5) - 1)
+        else:
+            loss = curl_JxB_loss
+        loss =  loss / normalization
+        #
+        return loss.mean()
 
 class DivergenceLoss(BaseLoss):
 
@@ -90,18 +142,18 @@ def __init__(self, base_radius, Mm_per_ds, **kwargs):
 
     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.norm(coords, dim=-1)
         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 = torch.norm(torch.cross(b, coords, dim=-1), dim=-1)
         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):
@@ -132,7 +184,8 @@ def forward(self, jac_matrix, coords, *args, **kwargs):
 
         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
+            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()
@@ -171,15 +224,16 @@ def forward(self, b, jac_matrix, coords, *args, **kwargs):
                 (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
+        radius_weight = torch.norm(coords, dim=-1)
+        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 = torch.relu(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):
@@ -197,6 +251,7 @@ def forward(self, b, coords, *args, **kwargs):
 
         return energy_regularization
 
+
 class NaNLoss(BaseLoss):
 
     def forward(self, boundary_b, b_true):
@@ -207,7 +262,24 @@ def forward(self, boundary_b, b_true):
         return min_energy_NaNs_regularization
 
 
-class LosTrvAziBoundaryLoss(BaseLoss):
+class LosTrvBoundaryLoss(BaseLoss):
+
+    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)
+
+        b_los_trv_pred = b_pred[..., :2]
+        b_los_trv_true = b_true[..., :2]
+
+        # compute diff
+        b_diff = b_los_trv_pred - b_los_trv_true
+        b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1))
+
+        return b_diff
+
+
+class AziBoundaryLoss(BaseLoss):
 
     def __init__(self, disambiguate=True, **kwargs):
         super().__init__(**kwargs)
@@ -218,16 +290,64 @@ def forward(self, b, b_true, transform=None, *args, **kwargs):
         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)
+        b_azi_true = b_true[..., 2] % torch.pi
+        b_azi_pred = b_pred[..., 2] % torch.pi
+        b_diff = (b_azi_pred - b_azi_true).pow(2) * b_true[..., 1]  # weight by transverse field
+        b_diff = b_diff.mean()
+
+        return b_diff
+
+
+class LosTrvAziBoundaryLoss(BaseLoss):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def forward(self, b, b_true, flip=None, transform=None, *args, **kwargs):
+        # apply transforms
+        bxyz_pred = torch.einsum('ijk,ik->ij', transform, b) if transform is not None else b
+        bxyz_true = los_trv_azi_to_img(b_true, f=torch)
+
+        if flip is not None:
+            # compute flipped B_xyz
+            flipped_azi = (b_true[..., 2:3] + torch.pi)
+            blta_flipped_true = torch.cat([b_true[..., :2], flipped_azi], dim=-1)
+            bxyz_flipped_true = los_trv_azi_to_img(blta_flipped_true, f=torch)
+
+            # compute diff for both cases
+            b_diff = (bxyz_pred - bxyz_true).pow(2).sum(-1)
+            b_diff_flipped = (bxyz_pred - bxyz_flipped_true).pow(2).sum(-1)
+
+            # weighted loss
+            flip = flip[..., 0]
+            loss = b_diff * (1 - flip) + b_diff_flipped * flip
+
+            return loss.mean()
 
         # compute diff
         b_diff = bxyz_pred - bxyz_true
-        b_diff = torch.mean(torch.nansum(b_diff.pow(2), -1))
+        b_diff = b_diff.pow(2).sum(-1).mean()
 
         return b_diff
 
 
+class LosBoundaryLoss(BaseLoss):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def forward(self, b, b_true, transform=None, *args, **kwargs):
+        # apply transforms
+        bxyz_pred = torch.einsum('ijk,ik->ij', transform, b) if transform is not None else b
+
+        b_los_pred = bxyz_pred[..., 2]
+        b_los_true = b_true[..., 0]
+
+        # compute diff
+        b_diff = (b_los_pred - b_los_true).pow(2).mean()
+        return b_diff
+
+
 class BoundaryLoss(BaseLoss):
 
     def forward(self, b, b_true, transform=None, b_err=None, *args, **kwargs):
@@ -252,24 +372,32 @@ def forward(self, coords, original_coords, height_range, *args, **kwargs):
         return height_regularization
 
 
+class AzimuthDisambiguationLoss(BaseLoss):
+
+    def __init__(self, power=4.0, **kwargs):
+        super().__init__(**kwargs)
+        self.power = power
+
+    def forward(self, flip, *args, **kwargs):
+        loss = (flip - 0.5).abs().pow(self.power) / 0.5 ** self.power
+        return loss.mean()
+
+
 class MinHeightLoss(BaseLoss):
 
-    def forward(self, b_true, coords, *args, **kwargs):
+    def forward(self, 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
+# mapping
+loss_module_mapping = {'boundary': BoundaryLoss, 'boundary_los_trv': LosTrvBoundaryLoss,
+                       'boundary_azi': AziBoundaryLoss,
+                       'boundary_los_trv_azi': LosTrvAziBoundaryLoss, 'boundary_los': LosBoundaryLoss,
+                       'divergence': DivergenceLoss, 'force_free': ForceFreeLoss, 'potential': PotentialLoss,
+                       'height': HeightLoss, 'NaNs': NaNLoss, 'radial': RadialLoss,
+                       'min_height': MinHeightLoss, 'energy_gradient': EnergyGradientLoss, 'energy': EnergyLoss,
+                       'magneto_static': MagnetoStaticLoss, 'implicit_magnetostatic': ImplicitMagnetoStaticLoss,
+                       'azimuth_disambiguation': AzimuthDisambiguationLoss}
diff --git a/nf2/train/mapping.py b/nf2/train/mapping.py
index f6936e3..a2fd06b 100644
--- a/nf2/train/mapping.py
+++ b/nf2/train/mapping.py
@@ -3,15 +3,24 @@
 from nf2.loader.fits import FITSDataset
 from nf2.loader.spherical import SphericalSliceDataset
 from nf2.train.callback import SphericalSlicesCallback, SlicesCallback, MetricsCallback, BoundaryCallback, \
-    LosTrvAziBoundaryCallback
+    LosTrvAziBoundaryCallback, DisambiguationCallback
 
 
-def load_callbacks(data_module):
+def load_callbacks(data_module, additional_callbacks=[]):
     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]
+        for callback_config in additional_callbacks:
+            ds_id = callback_config['ds_id']
+            callback_type = callback_config['type']
+            if ds_id == validation_dataset_key:
+                if callback_type == 'disambiguation':
+                    callback = DisambiguationCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
+                else:
+                    raise NotImplementedError(f'Callback type {callback_type} not implemented')
+                callbacks += [callback]
         if isinstance(ds, SphereSlicesDataset):
             callback = SphericalSlicesCallback(validation_dataset_key, ds.cube_shape, G_per_dB, Mm_per_ds)
         elif isinstance(ds, SlicesDataset):
diff --git a/nf2/train/model.py b/nf2/train/model.py
index 31348cb..274b066 100644
--- a/nf2/train/model.py
+++ b/nf2/train/model.py
@@ -19,16 +19,6 @@ 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):
     def __init__(self, w0=1.):
         super().__init__()
@@ -38,59 +28,14 @@ def forward(self, x):
         return torch.sin(self.w0 * x)
 
 
-class HeightTransformModel(nn.Module):
-
-    def __init__(self, in_coords, ds_id, validation_ds_id=[], dim=256, positional_encoding=True, ):
-        super().__init__()
-        if positional_encoding:
-            self.posenc = PositionalEncoding(20, in_coords)
-            d_in = nn.Linear(in_coords * 40, dim)
-            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 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]
-        output_coords = torch.cat([coords[:, :2], z_coords], -1)
-        return output_coords
-
-
 class RadialTransformModel(nn.Module):
 
     def __init__(self, in_coords, dim, positional_encoding=True, ds_ids=[]):
         super().__init__()
         if positional_encoding:
-            posenc = PositionalEncoding(20, in_coords)
-            d_in = nn.Linear(in_coords * 40, dim)
+            posenc = GaussianPositionalEncoding(num_freqs=20,
+                                                d_input=in_coords)
+            d_in = nn.Linear(posenc.d_output, dim)
             self.d_in = nn.Sequential(posenc, d_in)
         else:
             self.d_in = nn.Linear(in_coords, dim)
@@ -124,7 +69,7 @@ def transform(self, coords, obs_coords, height_range, **kwargs):
 
 class GenericModel(nn.Module):
 
-    def __init__(self, in_coords, out_coords, dim=256, encoding=None, activation='swish'):
+    def __init__(self, in_coords, out_coords, dim=256, n_layers=8, encoding=None, activation='sine'):
         super().__init__()
         if encoding is None or encoding == 'none':
             self.d_in = nn.Linear(in_coords, dim)
@@ -132,18 +77,21 @@ def __init__(self, in_coords, out_coords, dim=256, encoding=None, activation='sw
             posenc = PositionalEncoding(20, in_coords)
             d_in = nn.Linear(in_coords * 40, dim)
             self.d_in = nn.Sequential(posenc, d_in)
+        elif encoding == 'gaussian':
+            posenc = GaussianPositionalEncoding(20, in_coords)
+            d_in = nn.Linear(posenc.d_output, dim)
+            self.d_in = nn.Sequential(posenc, d_in)
         else:
             raise NotImplementedError(f'Unknown encoding {encoding}')
-        lin = [nn.Linear(dim, dim) for _ in range(8)]
+        lin = [nn.Linear(dim, dim) for _ in range(n_layers)]
         self.linear_layers = nn.ModuleList(lin)
         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)])
+        self.activations = nn.ModuleList([activation_f() for _ in range(n_layers)])
 
     def forward(self, 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))
@@ -245,17 +193,32 @@ def forward(self, coords, compute_jacobian=True):
 
 class PositionalEncoding(nn.Module):
 
-    def __init__(self, num_freqs, in_features):
+    def __init__(self, num_freqs, in_features, max_freq=8):
         super().__init__()
-        frequencies = torch.randn(num_freqs, in_features)
-        self.frequencies = nn.Parameter(frequencies[None], requires_grad=True)
+        frequencies = 2 ** torch.linspace(0, max_freq, num_freqs)
+        self.frequencies = nn.Parameter(frequencies[None, :, None], requires_grad=False)
+        self.d_output = in_features * (num_freqs * 2)
 
     def forward(self, x):
-        encoded = x[:, None, :] * torch.pi * 2 ** self.frequencies
+        encoded = x[:, None, :] * torch.pi * self.frequencies
         encoded = encoded.reshape(x.shape[0], -1)
         encoded = torch.cat([torch.sin(encoded), torch.cos(encoded)], -1)
         return encoded
 
+class GaussianPositionalEncoding(nn.Module):
+
+    def __init__(self, num_freqs, d_input, scale=1.):
+        super().__init__()
+        frequencies = torch.randn(num_freqs, d_input) * scale
+        self.frequencies = nn.Parameter(frequencies[None], requires_grad=False)
+        self.d_output = d_input * (num_freqs * 2 + 1)
+
+    def forward(self, x):
+        encoded = x[:, None, :] * self.frequencies
+        encoded = encoded.reshape(x.shape[0], -1)
+        encoded = torch.cat([x, torch.sin(encoded), torch.cos(encoded)], -1)
+        return encoded
+
 
 class ObserverTransformer(nn.Module):
 
@@ -304,6 +267,7 @@ def calculate_current(b, coords, jac_matrix=None):
     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]
diff --git a/nf2/train/module.py b/nf2/train/module.py
index 93ed642..597f2c6 100644
--- a/nf2/train/module.py
+++ b/nf2/train/module.py
@@ -1,19 +1,21 @@
 import copy
 import logging
 
+import torch
 from pytorch_lightning import LightningModule
+from torch import nn
 from torch.optim.lr_scheduler import ExponentialLR
 
-from nf2.train.loss import *
-from nf2.train.model import BModel, VectorPotentialModel, HeightTransformModel, MagnetoStaticModel, MagnetoStaticModelV2
+from nf2.train.loss import loss_module_mapping
+from nf2.train.model import BModel, VectorPotentialModel, MagnetoStaticModel, MagnetoStaticModelV2
+from nf2.train.transform import HeightTransformModel, AzimuthTransformModel
 
 
 class NF2Module(LightningModule):
 
     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):
+                 transforms=[], meta_path=None, **kwargs):
         """
         The main module for training the neural field model.
 
@@ -44,14 +46,8 @@ def __init__(self, validation_mapping, data_config, model_kwargs=None, loss_conf
             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:
-            raise ValueError(f"Invalid coordinate mapping model: {transform_type}")
+        transform_modules = self.load_transfrom_config(transforms)
+        self.transform_modules = nn.ModuleList(transform_modules)
 
         # load meta state
         if meta_path:
@@ -70,13 +66,36 @@ def __init__(self, validation_mapping, data_config, model_kwargs=None, loss_conf
                 {"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
+        loss_modules, lambdas, scheduled_lambdas = self.load_loss_config(data_config, loss_config)
+        self.loss_modules = nn.ModuleDict(loss_modules)
+        self.scheduled_lambdas = nn.ParameterDict(scheduled_lambdas)
+        self.lambdas = nn.ParameterDict(lambdas)
+
+        self.model = model
+        self.validation_mapping = validation_mapping
+        self.lr_params = lr_params
+
+        #
+        self.validation_outputs = {}
+
+    def load_transfrom_config(self, transforms):
+        transform_modules = []
+        for transform_config in transforms:
+            transform_config = copy.deepcopy(transform_config)
+
+            transform_type = transform_config.pop('type')
+            ds_ids = transform_config.pop('ds_id')
+            if transform_type == 'height':
+                transform_module = HeightTransformModel(**transform_config, ds_id=ds_ids)
+            elif transform_type == 'azimuth':
+                transform_module = AzimuthTransformModel(**transform_config, ds_id=ds_ids)
+            else:
+                raise ValueError(f"Invalid coordinate mapping model: {transform_type}")
+            transform_modules.append(transform_module)
+        return transform_modules
+
+    def load_loss_config(self, data_config, loss_config):
         scheduled_lambdas = {}
         lambdas = {}
         loss_modules = {}
@@ -111,21 +130,12 @@ def __init__(self, validation_mapping, data_config, model_kwargs=None, loss_conf
             # 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.validation_outputs = {}
-        self.loss_modules = nn.ModuleDict(loss_modules)
+        return loss_modules, lambdas, scheduled_lambdas
 
     def configure_optimizers(self):
         parameters = list(self.model.parameters())
-        if self.transform_module is not None:
-            parameters += list(self.transform_module.parameters())
+        parameters += list(self.transform_modules.parameters())
+        parameters += list(self.loss_modules.parameters())
         if isinstance(self.lr_params, dict):
             lr_start = self.lr_params['start']
             lr_end = self.lr_params['end']
@@ -150,8 +160,7 @@ def training_step(self, batch, batch_nb):
             batch[k]['coords'].requires_grad = True
 
         # transform batch
-        if self.transform_module is not None:
-            batch = self.transform_module(batch)
+        self.apply_transforms(batch)
 
         # concatenate all points
         coords = torch.cat([batch[k]['coords'] for k in loader_keys], 0)
@@ -175,20 +184,45 @@ def training_step(self, batch, batch_nb):
 
         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)
+            try:
+                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)
+            except Exception as e:
+                logging.error(f"Error in loss module: {name}")
+                raise e
         total_loss = sum([self.lambdas[k] * loss_dict[k] for k in loss_dict.keys()])
 
         loss_dict['loss'] = total_loss
         return loss_dict
 
+    def apply_transforms(self, batch):
+        loader_keys = list(batch.keys())
+        for transform_module in self.transform_modules:
+            ds_ids = transform_module.ds_id
+            tensor_ids = transform_module.tensor_ids
+            # concatenate all tensors
+            transform_ds_ids = [k for k in loader_keys if k in ds_ids]
+            transform_batch = {tensor_id: torch.cat([batch[k][tensor_id] for k in transform_ds_ids], 0) for tensor_id in
+                               tensor_ids}
+            batch_lengths = [batch[k][tensor_ids[0]].shape[0] for k in transform_ds_ids]
+            transformed_batch = transform_module(transform_batch)
+
+            # update batch
+            for k in transformed_batch.keys():
+                value = transformed_batch[k]
+
+                number_idx = 0
+                for ds_id, batch_length in zip(transform_ds_ids, batch_lengths):
+                    batch[ds_id][k] = value[number_idx:number_idx + batch_length]
+                    number_idx += batch_length
+
     @torch.no_grad()
     def on_train_batch_end(self, outputs, batch, batch_idx) -> None:
         # update lambda parameters and log
@@ -221,10 +255,14 @@ def validation_step(self, batch, batch_nb, dataloader_idx):
         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)
+        for transform_module in self.transform_modules:
+            ds_ids = transform_module.ds_id
+            if self.validation_mapping[dataloader_idx] in ds_ids:
+                transform_out = transform_module(batch)
+                for k, v in transform_out.items():
+                    batch[k] = v
 
+        coords = batch['coords']
         model_out = self.model(coords)
 
         jac_matrix = model_out['jac_matrix']
@@ -295,7 +333,6 @@ def on_load_checkpoint(self, checkpoint):
 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
+                  'data': data_module.config,
+                  'transforms': nf2.transform_modules, }
     torch.save(save_state, save_path)
diff --git a/nf2/train/transform.py b/nf2/train/transform.py
new file mode 100644
index 0000000..825c61d
--- /dev/null
+++ b/nf2/train/transform.py
@@ -0,0 +1,44 @@
+import torch
+from torch import nn
+
+from nf2.train.model import GenericModel
+
+
+class BaseTransformModel(nn.Module):
+
+    def __init__(self, ds_id, tensor_ids):
+        super().__init__()
+        self.ds_id = ds_id if isinstance(ds_id, list) else [ds_id]
+        self.tensor_ids = tensor_ids if isinstance(tensor_ids, list) else [tensor_ids]
+
+
+class AzimuthTransformModel(BaseTransformModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(tensor_ids=['coords', 'b_true'], **kwargs)
+        self.model = GenericModel(in_coords=3, out_coords=1, n_layers=8, dim=128, encoding='gaussian')
+
+    def forward(self, batch, **kwargs):
+        coords = batch['coords']
+
+        # flip probability between 0 and 1
+        flip = torch.sigmoid(self.model(coords))
+
+        return {'flip': flip}
+
+
+class HeightTransformModel(BaseTransformModel):
+
+    def __init__(self, **kwargs):
+        super().__init__(tensor_ids=['coords', 'height_range', 'b_true'], **kwargs)
+        self.mapping_module = GenericModel(in_coords=3, out_coords=1, n_layers=4, dim=128, encoding='gaussian')
+
+    def forward(self, batch):
+        coords = batch['coords']
+        height_range = batch['height_range']
+
+        z_coords = self.mapping_module(coords)
+        z_coords = torch.sigmoid(z_coords) * (height_range[:, 1:2] - height_range[:, 0:1]) + height_range[:, 0:1]
+        transformed_coords = torch.cat([coords[:, :2], z_coords], -1)
+
+        return {'coords': transformed_coords, 'original_coords': coords}
diff --git a/nf2/train/util.py b/nf2/train/util.py
new file mode 100644
index 0000000..51546ed
--- /dev/null
+++ b/nf2/train/util.py
@@ -0,0 +1,14 @@
+import yaml
+
+
+def load_yaml_config(yaml_config_file, overwrite_args=None):
+    overwrite_args = [] if overwrite_args is None else overwrite_args
+    assert all([k.startswith('--') for k in overwrite_args[::2]]), \
+        'Only accept --config and overwrite arguments (must start with --)'
+    overwrite_args = {k.replace('--', ''): v for k, v in zip(overwrite_args[::2], overwrite_args[1::2])}
+    with open(yaml_config_file) as f:
+        config_str = f.read()
+    for overwrite_key, overwrite_value in overwrite_args.items():
+        config_str = config_str.replace('{%s}' % overwrite_key, overwrite_value)
+    config = yaml.safe_load(config_str)
+    return config


Test patch:
diff --git a/config/fd_2173_potential_test.json b/config/fd_2173_potential_test.json
deleted file mode 100644
index c79e9ad..0000000
--- a/config/fd_2173_potential_test.json
+++ /dev/null
@@ -1,60 +0,0 @@
-{
-  "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