From f6465ec47b324af214e8695aa5a1297f322bb2a9 Mon Sep 17 00:00:00 2001
From: Didier Vezinet <dvezinet@cfs.energy>
Date: Sat, 26 Oct 2024 12:03:08 -0400
Subject: [PATCH 1/3] [#981] diag loading from json more robust vs missing los
 and vos

---
 tofu/data/_class08_loadfromjson.py | 123 ++++++++++++++++++-----------
 1 file changed, 75 insertions(+), 48 deletions(-)

diff --git a/tofu/data/_class08_loadfromjson.py b/tofu/data/_class08_loadfromjson.py
index 0ba83619a..6c5eecd7b 100644
--- a/tofu/data/_class08_loadfromjson.py
+++ b/tofu/data/_class08_loadfromjson.py
@@ -266,29 +266,30 @@ def _add_diagnostic(
     for kcam in lcam:
 
         # add ray
-        klos = doptics[kcam]['los_key']
-        coll.add_rays(
-            key=klos,
-            # start
-            start_x=doptics[kcam]['los_x_start']['data'],
-            start_y=doptics[kcam]['los_y_start']['data'],
-            start_z=doptics[kcam]['los_z_start']['data'],
-            # pts
-            pts_x=doptics[kcam]['los_x_end']['data'],
-            pts_y=doptics[kcam]['los_y_end']['data'],
-            pts_z=doptics[kcam]['los_z_end']['data'],
-            # angles
-            alpha=doptics[kcam]['los_alpha']['data'],
-            dalpha=doptics[kcam]['los_dalpha']['data'],
-            dbeta=doptics[kcam]['los_dbeta']['data'],
-        )
+        klos = doptics[kcam].get('los_key')
+        if klos is not None:
+            coll.add_rays(
+                key=klos,
+                # start
+                start_x=doptics[kcam]['los_x_start']['data'],
+                start_y=doptics[kcam]['los_y_start']['data'],
+                start_z=doptics[kcam]['los_z_start']['data'],
+                # pts
+                pts_x=doptics[kcam]['los_x_end']['data'],
+                pts_y=doptics[kcam]['los_y_end']['data'],
+                pts_z=doptics[kcam]['los_z_end']['data'],
+                # angles
+                alpha=doptics[kcam]['los_alpha']['data'],
+                dalpha=doptics[kcam]['los_dalpha']['data'],
+                dbeta=doptics[kcam]['los_dbeta']['data'],
+            )
 
-        # adjust ref to match camera
-        ref = tuple(
-            [coll.dobj['rays'][klos]['ref'][0]]
-            + list(coll.dobj['camera'][kcam]['dgeom']['ref'])
-        )
-        coll._dobj['rays'][klos]['ref'] = ref
+            # adjust ref to match camera
+            ref = tuple(
+                [coll.dobj['rays'][klos]['ref'][0]]
+                + list(coll.dobj['camera'][kcam]['dgeom']['ref'])
+            )
+            coll._dobj['rays'][klos]['ref'] = ref
 
         # store in diag
         coll._dobj['diagnostic'][din['key']]['doptics'][kcam]['los'] = klos
@@ -299,24 +300,28 @@ def _add_diagnostic(
     for kcam in lcam:
 
         # add data
-        ketend = doptics[kcam]['etendue_key']
-        ref = coll.dobj['camera'][kcam]['dgeom']['ref']
-        coll.add_data(
-            key=ketend,
-            data=doptics[kcam]['etendue']['data'],
-            units=doptics[kcam]['etendue']['units'],
-            ref=ref,
-        )
+        ketend = doptics[kcam].get('etendue_key')
+        etend_type = doptics[kcam].get('etend_type')
+        if ketend is not None:
+            ref = coll.dobj['camera'][kcam]['dgeom']['ref']
+            coll.add_data(
+                key=ketend,
+                data=doptics[kcam]['etendue']['data'],
+                units=doptics[kcam]['etendue']['units'],
+                ref=ref,
+            )
 
         # store in diag
         coll._dobj['diagnostic'][din['key']]['doptics'][kcam]['etendue'] = ketend
-        coll._dobj['diagnostic'][din['key']]['doptics'][kcam]['etend_type'] = doptics[kcam]['etend_type']
+        coll._dobj['diagnostic'][din['key']]['doptics'][kcam]['etend_type'] = etend_type
 
     # -----------
     # vos
+    # -----------
 
     for kcam in lcam:
 
+        # --------
         # add data
 
         for kp in ['pcross_x0', 'pcross_x1', 'phor_x0', 'phor_x1']:
@@ -325,27 +330,49 @@ def _add_diagnostic(
             ref = doptics[kcam][kp]['ref']
 
             # add ref if needed
-            lr = [rr for rr in ref if rr not in coll.dref.keys()]
-            for rr in lr:
-                ii = ref.index(rr)
-                coll.add_ref(
-                    key=rr,
-                    size=np.array(doptics[kcam][kp]['data']).shape[ii],
-                )
+            if ref is not None:
+                lr = [rr for rr in ref if rr not in coll.dref.keys()]
+                for rr in lr:
+                    ii = ref.index(rr)
+                    coll.add_ref(
+                        key=rr,
+                        size=np.array(doptics[kcam][kp]['data']).shape[ii],
+                    )
 
             # add data
-            coll.add_data(
-                key=kpi,
-                data=doptics[kcam][kp]['data'],
-                units=doptics[kcam][kp]['units'],
-                ref=tuple(ref),
-            )
+            if kpi is not None:
+                coll.add_data(
+                    key=kpi,
+                    data=doptics[kcam][kp]['data'],
+                    units=doptics[kcam][kp]['units'],
+                    ref=tuple(ref),
+                )
+
+        # -----------------
+        # store vos in diag
+
+        pcross0 = doptics[kcam].get('pcross_x0', {}).get('key')
+        pcross1 = doptics[kcam].get('pcross_x1', {}).get('key')
+        phor0 = doptics[kcam].get('phor_x0', {}).get('key')
+        phor1 = doptics[kcam].get('phor_x1', {}).get('key')
+        dphi = doptics[kcam].get('dphi', {}).get('data')
+
+        if pcross0 is not None:
+            pcross = (pcross0, pcross1)
+            phor = (phor0, phor1)
+            dphi = np.array(dphi)
+        else:
+            pcross = None
+            phor = None
+            dphi = None
+
+        # -----------------
+        # store phor in diag
 
-        # store in diag
         coll._dobj['diagnostic'][din['key']]['doptics'][kcam]['dvos'] = {
-            'pcross': (doptics[kcam]['pcross_x0']['key'], doptics[kcam]['pcross_x1']['key']),
-            'phor': (doptics[kcam]['phor_x0']['key'], doptics[kcam]['phor_x1']['key']),
-            'dphi': np.array(doptics[kcam]['dphi']['data']),
+            'pcross': pcross,
+            'phor': phor,
+            'dphi': dphi,
         }
 
     return
\ No newline at end of file

From 48a4b55f7fedc7016d2b20649986661e15ad366d Mon Sep 17 00:00:00 2001
From: Didier Vezinet <dvezinet@cfs.energy>
Date: Sat, 26 Oct 2024 12:05:32 -0400
Subject: [PATCH 2/3] [#981] saving diag to json more robust to missing los and
 vos

---
 tofu/data/_class08_save2json.py | 166 ++++++++++++++++----------------
 1 file changed, 85 insertions(+), 81 deletions(-)

diff --git a/tofu/data/_class08_save2json.py b/tofu/data/_class08_save2json.py
index 10c27d451..390ca04cd 100644
--- a/tofu/data/_class08_save2json.py
+++ b/tofu/data/_class08_save2json.py
@@ -220,7 +220,7 @@ def _extract_diagnostic(
 
         for k0 in ['etendue', 'amin', 'amax']:
 
-            kd = doptics[kcam][k0]
+            kd = doptics[kcam].get(k0)
             if kd is None:
                 continue
             dout['diagnostic']['doptics'][kcam][k0] = {
@@ -236,98 +236,102 @@ def _extract_diagnostic(
         # ---------
         # los
 
-        klos = doptics[kcam]['los']
-        ptsx, ptsy, ptsz = coll.get_rays_pts(klos)
-
-        klos = doptics[kcam]['los']
-
-        dout['diagnostic']['doptics'][kcam].update({
-            'los_x_start': {
-                'data': ptsx[0, ...],
-                'units': 'm',
-            },
-            'los_y_start': {
-                'data': ptsy[0, ...],
-                'units': 'm',
-            },
-            'los_z_start': {
-                'data': ptsz[0, ...],
-                'units': 'm',
-            },
-            'los_x_end': {
-                'data': ptsx[1:, ...],
-                'units': 'm',
-            },
-            'los_y_end': {
-                'data': ptsy[1:, ...],
-                'units': 'm',
-            },
-            'los_z_end': {
-                'data': ptsz[1:, ...],
-                'units': 'm',
-            },
-            'los_key': klos,
-            'los_alpha': {
-                'data': coll.ddata[coll.dobj['rays'][klos]['alpha']]['data'],
-                'units': 'rad'
-            },
-            'los_dalpha': {
-                'data': coll.ddata[coll.dobj['rays'][klos]['reflect_dalpha']]['data'],
-                'units': 'rad'
-            },
-            'los_dbeta': {
-                'data': coll.ddata[coll.dobj['rays'][klos]['reflect_dbeta']]['data'],
-                'units': 'rad'
-            },
-        })
+        klos = doptics[kcam].get('los')
+        if klos is not None:
+            ptsx, ptsy, ptsz = coll.get_rays_pts(klos)
+
+            klos = doptics[kcam]['los']
+
+            dout['diagnostic']['doptics'][kcam].update({
+                'los_x_start': {
+                    'data': ptsx[0, ...],
+                    'units': 'm',
+                },
+                'los_y_start': {
+                    'data': ptsy[0, ...],
+                    'units': 'm',
+                },
+                'los_z_start': {
+                    'data': ptsz[0, ...],
+                    'units': 'm',
+                },
+                'los_x_end': {
+                    'data': ptsx[1:, ...],
+                    'units': 'm',
+                },
+                'los_y_end': {
+                    'data': ptsy[1:, ...],
+                    'units': 'm',
+                },
+                'los_z_end': {
+                    'data': ptsz[1:, ...],
+                    'units': 'm',
+                },
+                'los_key': klos,
+                'los_alpha': {
+                    'data': coll.ddata[coll.dobj['rays'][klos]['alpha']]['data'],
+                    'units': 'rad'
+                },
+                'los_dalpha': {
+                    'data': coll.ddata[coll.dobj['rays'][klos]['reflect_dalpha']]['data'],
+                    'units': 'rad'
+                },
+                'los_dbeta': {
+                    'data': coll.ddata[coll.dobj['rays'][klos]['reflect_dbeta']]['data'],
+                    'units': 'rad'
+                },
+            })
 
         # ------------
         # vos - pcross
 
-        pc0, pc1 = doptics[kcam]['dvos']['pcross']
-        dout['diagnostic']['doptics'][kcam].update({
-            'pcross_x0': {
-                'key': pc0,
-                'data': coll.ddata[pc0]['data'],
-                'units': coll.ddata[pc0]['units'],
-                'ref': coll.ddata[pc0]['ref'],
-            },
-            'pcross_x1': {
-                'key': pc1,
-                'data': coll.ddata[pc1]['data'],
-                'units': coll.ddata[pc1]['units'],
-                'ref': coll.ddata[pc1]['ref'],
-            },
-        })
+        if doptics[kcam].get('dvos', {}).get('pcross') is not None:
+            pc0, pc1 = doptics[kcam]['dvos']['pcross']
+            dout['diagnostic']['doptics'][kcam].update({
+                'pcross_x0': {
+                    'key': pc0,
+                    'data': coll.ddata[pc0]['data'],
+                    'units': coll.ddata[pc0]['units'],
+                    'ref': coll.ddata[pc0]['ref'],
+                },
+                'pcross_x1': {
+                    'key': pc1,
+                    'data': coll.ddata[pc1]['data'],
+                    'units': coll.ddata[pc1]['units'],
+                    'ref': coll.ddata[pc1]['ref'],
+                },
+            })
 
         # ------------
         # vos - phor
 
-        ph0, ph1 = doptics[kcam]['dvos']['phor']
-        dout['diagnostic']['doptics'][kcam].update({
-            'phor_x0': {
-                'key': ph0,
-                'data': coll.ddata[ph0]['data'],
-                'units': coll.ddata[ph0]['units'],
-                'ref': coll.ddata[ph0]['ref'],
-            },
-            'phor_x1': {
-                'key': ph1,
-                'data': coll.ddata[ph1]['data'],
-                'units': coll.ddata[ph1]['units'],
-                'ref': coll.ddata[ph1]['ref'],
-            },
-        })
+        if doptics[kcam].get('dvos', {}).get('phor') is not None:
+            ph0, ph1 = doptics[kcam]['dvos']['phor']
+            dout['diagnostic']['doptics'][kcam].update({
+                'phor_x0': {
+                    'key': ph0,
+                    'data': coll.ddata[ph0]['data'],
+                    'units': coll.ddata[ph0]['units'],
+                    'ref': coll.ddata[ph0]['ref'],
+                },
+                'phor_x1': {
+                    'key': ph1,
+                    'data': coll.ddata[ph1]['data'],
+                    'units': coll.ddata[ph1]['units'],
+                    'ref': coll.ddata[ph1]['ref'],
+                },
+            })
 
         # ------------
         # vos - dphi
 
-        dout['diagnostic']['doptics'][kcam].update({
-            'dphi': {
-                'data': doptics[kcam]['dvos']['dphi'],
-                'units': 'rad',
-            },
-        })
+        if doptics[kcam].get('dvos', {}).get('dphi') is not None:
+            dout['diagnostic']['doptics'][kcam].update({
+                'dphi': {
+                    'data': doptics[kcam]['dvos']['dphi'],
+                    'units': 'rad',
+                },
+            })
 
     return dout, dcls
 

From ed7e080f35a7a03195a7fe99a6ece94578564652 Mon Sep 17 00:00:00 2001
From: Didier Vezinet <dvezinet@cfs.energy>
Date: Sat, 26 Oct 2024 12:49:07 -0400
Subject: [PATCH 3/3] [#981] add_diagnostic_from_file(returnas=bool)
 implemented, with proper default values

---
 tofu/data/_class08_Diagnostic.py         | 10 +++++++---
 tofu/data/_class08_loadfromjson.py       | 21 ++++++++++++++++++---
 tofu/data/_class08_saveload_from_file.py |  2 ++
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/tofu/data/_class08_Diagnostic.py b/tofu/data/_class08_Diagnostic.py
index 89788db8c..a8bc488e0 100644
--- a/tofu/data/_class08_Diagnostic.py
+++ b/tofu/data/_class08_Diagnostic.py
@@ -1478,22 +1478,26 @@ def save_diagnostic_to_file(
     def add_diagnostic_from_file(
         self,
         pfe=None,
+        returnas=False,
     ):
         """ Adds a diagnostic instance (and necessary optics) from json file
 
         Parameters
         ----------
-        pfe : str
+        pfe :     str
             path/file.ext to desired file
+        returnas: bool
+            whether to return the Collection instance
 
         Returns
         -------
-        TYPE
-            DESCRIPTION.
+        coll:    Collection instance
+            optional
 
         """
 
         return _saveload_from_file.load(
             coll=self,
             pfe=pfe,
+            returnas=returnas,
         )
\ No newline at end of file
diff --git a/tofu/data/_class08_loadfromjson.py b/tofu/data/_class08_loadfromjson.py
index 6c5eecd7b..6dfb901d4 100644
--- a/tofu/data/_class08_loadfromjson.py
+++ b/tofu/data/_class08_loadfromjson.py
@@ -10,6 +10,7 @@
 
 
 import numpy as np
+import datastock as ds
 
 
 # #################################################################
@@ -24,15 +25,17 @@ def main(
     pfe=None,
     coll=None,
     fname=None,
+    returnas=None,
 ):
 
     # ----------------
     # check inputs
     # --------------
 
-    dout = _check(
+    dout, returnas = _check(
         coll=coll,
         pfe=pfe,
+        returnas=returnas,
     )
 
     # ----------------
@@ -63,7 +66,8 @@ def main(
         din=dout['diagnostic'],
     )
 
-    return coll
+    if returnas is True:
+        return coll
 
 
 # #################################################################
@@ -75,6 +79,7 @@ def main(
 def _check(
     coll=None,
     pfe=None,
+    returnas=None,
 ):
 
     # --------------------
@@ -99,7 +104,17 @@ def _check(
         )
         raise Exception(msg)
 
-    return dout
+    # -------------
+    # returnas
+    # -------------
+
+    returnas = ds._generic_check._check_var(
+        returnas, 'returnas',
+        types=bool,
+        default=True,
+    )
+
+    return dout, returnas
 
 
 # #################################################################
diff --git a/tofu/data/_class08_saveload_from_file.py b/tofu/data/_class08_saveload_from_file.py
index 1357a1dd2..959fc84b6 100644
--- a/tofu/data/_class08_saveload_from_file.py
+++ b/tofu/data/_class08_saveload_from_file.py
@@ -118,6 +118,7 @@ def save(
 def load(
     pfe=None,
     coll=None,
+    returnas=None,
 ):
 
     # ----------------
@@ -138,6 +139,7 @@ def load(
             coll=coll,
             pfe=pfe,
             fname=fname,
+            returnas=returnas,
         )