Jlorincz/common testing

.github/workflows/tests.yml

@@ -86,6 +86,8 @@ jobs:
           echo "${{ secrets.CMOD_LOGIN }}" \
           | tee ~/logbook.sybase_login \
           | sha256sum
+          echo DISPY_TOKAMAK=CMOD \
+          | tee -a "$GITHUB_ENV"
 
       - name: Setup DIII-D
         if: ${{ matrix.tokamak == 'DIII-D' }}
@@ -107,7 +109,7 @@ jobs:
           | sha256sum
           echo -e "[FreeTDS]\nDescription = FreeTDS\nDriver = $TDS" \
           | sudo tee -a /etc/odbcinst.ini
-          echo DIIID_TEST=1 \
+          echo DISPY_TOKAMAK=D3D \
           | tee -a "$GITHUB_ENV"
 
       - name: Setup Python
@@ -139,12 +141,8 @@ jobs:
       - name: Test EFIT
         run: python examples/efit.py
 
-      - name: Test quick
-        run: pytest -v tests/test_quick.py
-
-      - name: Test features
-        if: ${{ matrix.tokamak == 'C-MOD' }}
-        run: pytest -v --durations=0 tests/test_cmod_features.py
+      - name: Run all tests
+        run: pytest -v --durations=0 tests
 
       - name: Close tunnel
         run: xargs -a ssh.pid kill -TERM

disruption_py/cli/evaluate_methods.py

@@ -1,141 +1,33 @@
 import argparse
 from contextlib import contextmanager
-from typing import Dict, List
 import numpy as np
-import pandas as pd
-from pandas.api.types import is_numeric_dtype
 import logging
-from disruption_py.handlers.cmod_handler import CModHandler
-from disruption_py.settings.log_settings import LogSettings
 from disruption_py.settings.shot_ids_request import ShotIdsRequestParams, shot_ids_request_runner
-from disruption_py.settings.shot_settings import ShotSettings
-from disruption_py.utils.constants import TIME_CONST
+from disruption_py.utils.eval.eval_against_sql import eval_against_sql, get_failure_statistics_string
 from disruption_py.utils.mappings.mappings_helpers import map_string_to_enum
 from disruption_py.utils.mappings.tokamak import Tokamak
-from disruption_py.utils.mappings.tokamak_helpers import get_tokamak_from_environment
-from disruption_py.utils.math_utils import matlab_gradient_1d_vectorized
+from disruption_py.utils.mappings.tokamak_helpers import get_tokamak_from_environment, get_tokamak_test_expected_failure_columns, get_tokamak_handler, get_tokamak_test_shot_ids
 
-CMOD_TEST_SHOTS = [
-    1150805012,   # Flattop Disruption
-    1150805013,     # No Disruption
-    1150805014,     # No Disruption
-    1150805015,     # Rampdown Disruption
-    1150805016,     # Rampdown Disruption
-    1150805017,     # Rampdown Disruption
-    1150805019,     # Rampdown Disruption
-    1150805020,     # Rampdown Disruption
-    1150805021,     # Rampdown Disruption
-    1150805022      # Flattop Disruption
-]
 
-TIME_EPSILON = 0.05 # Tolerance for taking the difference between two times [s]
-IP_EPSILON = 1e5    # Tolerance for taking the difference between two ip values [A]
-
-VAL_TOLERANCE = 0.01   # Tolerance for comparing values between MDSplus and SQL
-MATCH_FRACTION = 0.95   # Fraction of signals that must match between MDSplus and SQL
-
-def get_mdsplus_data(handler, shot_list):
-    shot_settings = ShotSettings(
-        efit_tree_name="efit18",
-        set_times_request="efit",
-        log_settings=LogSettings(
-            console_log_level=logging.ERROR
-        )
-    )
-    return handler.get_shots_data(
-        shot_ids_request=shot_list, 
-        shot_settings=shot_settings,
-        output_type_request="dict",
-    )
-
-def get_sql_data(handler, mdsplus_data : Dict, shot_list):
-    shot_data = {}
-    for shot_id in shot_list:
-        times = mdsplus_data[shot_id]['time']
-        sql_data =  handler.database.get_shots_data([shot_id])
-        shot_data[shot_id] = pd.merge_asof(times.to_frame(), sql_data, on='time', direction='nearest', tolerance=TIME_CONST)
-    return shot_data
-
-
-def test_data_match(sql_shot_df : pd.DataFrame, mdsplus_shot_df : pd.DataFrame, data_column : str):
-
-    sql_column_data = sql_shot_df[data_column].astype(np.float64)
-    mds_column_data = mdsplus_shot_df[data_column].astype(np.float64)
-
-    # copare data for numeric differences
-    relative_difference = np.where(
-        sql_column_data != 0, 
-        np.abs((mds_column_data - sql_column_data) / sql_column_data), 
-        np.where(mds_column_data != 0, np.inf, np.nan)
-    )
-    numeric_anomalies_mask = np.greater(relative_difference, VAL_TOLERANCE)
-
-    # compare data for nan differences
-    sql_is_nan_ = pd.isnull(sql_column_data)
-    mdsplus_is_nan = pd.isnull(mds_column_data)
-    nan_anomalies_mask = (sql_is_nan_ != mdsplus_is_nan)
-
-    anomalies = np.argwhere(numeric_anomalies_mask | nan_anomalies_mask)
-
-    return not (len(anomalies) / len(relative_difference) > 1 - MATCH_FRACTION)
-
-def evaluate_cmod_accuracy(shot_list : List = None):
-    """
-    Evaluate the accuracy of CMod methods.
-
-    Prints a short report on the methods that have suceeded and failed.
-    Success criteria is having more that 95% of results within 1% of known results. 
-    """
-    print("Evaluating accuracy of CMod methods...")
-    if shot_list is None or len(shot_list) == 0:
-        shot_list = CMOD_TEST_SHOTS
+def evaluate_accuracy(tokamak : Tokamak, shot_ids : list[int], fail_quick : bool = False, data_columns : list[str] = None):
+    handler = get_tokamak_handler(tokamak)
+    if shot_ids is None or len(shot_ids) == 0:
+        shot_ids = get_tokamak_test_shot_ids(tokamak)
     else:
-        shot_list = [int(shot_id) for shot_id in shot_list]
-
-    @contextmanager
-    def monkey_patch_numpy_gradient():
-        original_function = np.gradient
-        np.gradient = matlab_gradient_1d_vectorized
-        try:
-            yield
-        finally:
-            np.gradient = original_function
-
-    with monkey_patch_numpy_gradient():
-        return _evaluate_cmod_accuracy(shot_list)
-
-def _evaluate_cmod_accuracy(shot_list : List = None):
-    cmod_handler = CModHandler()
-    print("Getting data from MDSplus")
-    mdsplus_data = get_mdsplus_data(cmod_handler, shot_list)
-    print("Getting data from sql table")
-    sql_data = get_sql_data(cmod_handler, mdsplus_data, shot_list)
-
-    success_columns = set()
-    unknown_columns = set()
-    failure_columns = set()
-
-    for shot_id in shot_list:
-
-        mdsplus_shot_df : pd.DataFrame = mdsplus_data[shot_id]
-        sql_shot_df : pd.DataFrame = sql_data[shot_id]
+        shot_ids = [int(shot_id) for shot_id in shot_ids]
 
-        for data_column in mdsplus_shot_df.columns:
-            if data_column not in sql_shot_df.columns or not is_numeric_dtype(mdsplus_shot_df[data_column]):
-                unknown_columns.add(data_column)
-                continue
-                
-            if test_data_match(sql_shot_df, mdsplus_shot_df, data_column):
-                success_columns.add(data_column)
-            else:
-                failure_columns.add(data_column)
+    expected_failure_columns = get_tokamak_test_expected_failure_columns(tokamak)
+
+    data_differences = eval_against_sql(
+        handler=handler, 
+        shot_ids=shot_ids, 
+        expected_failure_columns=expected_failure_columns, 
+        fail_quick=fail_quick, 
+        test_columns=data_columns
+    )
 
-    success_columns = success_columns.difference(failure_columns)
-    unknown_columns = unknown_columns.difference(success_columns).difference(failure_columns)   
-    print(f"Successful Columns (failure criteria not met for any shot): {success_columns}") 
-    print(f"Columns with a failure: {failure_columns}") 
-    print(f"Columns that lacked testing data: {unknown_columns}") 
-    return success_columns, unknown_columns, failure_columns
+    print(get_failure_statistics_string(data_differences))
+
 
 def main(args):
     """
@@ -183,13 +75,15 @@ def main(args):
         )
         all_shot_ids = shot_ids_request_runner(args.shotlist, shot_ids_request_params)
 
-    if tokamak == Tokamak.CMOD:
-        evaluate_cmod_accuracy(all_shot_ids)
-    else:
-        print("Sorry, this tokamak is not currently supported.")
-
+    data_columns = [args.data_column] if args.data_column else None
+
+    print("Running evaluation...")
+    evaluate_accuracy(tokamak=tokamak, shot_ids=all_shot_ids, fail_quick=args.fail_quick, data_columns=data_columns)
 
+
 def get_parser():
     parser = argparse.ArgumentParser(description='Evaluate the accuracy of DisruptionPy methods on a Tokamak.')
     parser.add_argument('--shotlist', type=str, help='Path to file specifying a shotlist, leave blank for interactive mode', default=None)
+    parser.add_argument('--fail_quick', action='store_true', help='Fail quickly', default=False)
+    parser.add_argument('--data_column', type=str, help='Data column to test', default=None)
     return parser

disruption_py/handlers/__init__.py

@@ -1 +1 @@
-from .cmod_handler import CModHandler
+from .handler import Handler

disruption_py/mdsplus_integration/mds_connection.py

@@ -3,6 +3,8 @@
 import numpy as np
 import MDSplus
 
+from disruption_py.utils.utils import safe_cast
+
 class ProcessMDSConnection():
     """
     Abstract class for connecting to MDSplus.
@@ -142,7 +144,7 @@ def get_data(
 
         data = self.conn.get("_sig=" + path, arguments).data()
         if astype:
-            data = data.astype(astype, copy=False)
+            data = safe_cast(data, astype)
 
         return data
 
@@ -185,9 +187,9 @@ def get_data_with_dims(
         dims = [self.conn.get(f"dim_of(_sig,{dim_num})").data() for dim_num in dim_nums]
 
         if astype:
-            data = data.astype(astype, copy=False)
+            data = safe_cast(data, astype)
             if cast_all:
-                dims = [dim.astype(astype, copy=False) for dim in dims]
+                dims = [safe_cast(dim, astype) for dim in dims]
 
         return data, *dims
 
@@ -226,7 +228,7 @@ def get_dims(
         dims = [self.conn.get(f"dim_of({path},{d})").data() for d in dim_nums]
 
         if astype:
-            dims = [dim.astype(astype, copy=False) for dim in dims]
+            dims = [safe_cast(dim, astype) for dim in dims]
 
         return dims
 

disruption_py/settings/output_type_request.py

@@ -203,7 +203,7 @@ def get_results(self, params: FinishOutputTypeRequestParams):
         return self.results
 
     def stream_output_cleanup(self, params: FinishOutputTypeRequestParams):
-        self.results = []
+        self.results = {}
 
 class DataFrameOutputRequest(OutputTypeRequest):
     """
@@ -213,7 +213,8 @@ def __init__(self):
         self.results : pd.DataFrame = pd.DataFrame()
 
     def _output_shot(self, params : ResultOutputTypeRequestParams):
-        self.results = pd.concat([self.results, params.result], ignore_index=True)
+        if not params.result.empty and not params.result.isna().all().all():
+            self.results = pd.concat([self.results, params.result], ignore_index=True)
 
     def get_results(self, params: FinishOutputTypeRequestParams):
         return self.results

disruption_py/settings/set_times_request.py

@@ -272,6 +272,7 @@ def _get_times(self, params : SetTimesRequestParams) -> np.ndarray:
 _set_times_request_mappings: Dict[str, SetTimesRequest] = {
     "efit" : EfitSetTimesRequest(),
     "disruption" : DisruptionSetTimesRequest(),
+    "disruption_warning": {Tokamak.CMOD: EfitSetTimesRequest(), Tokamak.D3D: DisruptionSetTimesRequest()},
     "ip" : IpSetTimesRequest(),
 }
 # --8<-- [end:set_times_request_dict]

disruption_py/settings/shot_settings.py

@@ -53,7 +53,7 @@ class or in an included shot_data_request. All methods with at least one include
     set_times_request : SetTimesRequest
         The set times request to be used when setting the timebase for the shot. The retrieved data will
         be interpolated to this timebase. Can pass any SetTimesRequestType that resolves to a SetTimesRequest.
-        See SetTimesRequest for more details. Defaults to "efit".
+        See SetTimesRequest for more details. Defaults to "disruption_warning".
     signal_domain : SignalDomain
         The domain of the timebase that should be used when retrieving data for the shot. Either "full", 
         "flattop", or "rampup_and_flattop". Can pass either a SignalDomain or the associated string. Defaults 
@@ -91,7 +91,7 @@ class or in an included shot_data_request. All methods with at least one include
     shot_data_requests : List[ShotDataRequest] = field(default_factory=list)
 
     # Timebase setting
-    set_times_request : SetTimesRequest = "efit"
+    set_times_request : SetTimesRequest = "disruption_warning"
     signal_domain : SignalDomain = "full"
     use_existing_data_timebase : bool = False
     interpolation_method : InterpolationMethod = "linear"

disruption_py/shots/parameter_methods/cmod/basic_parameter_methods.py

@@ -6,7 +6,7 @@
 from disruption_py.settings.shot_data_request import ShotDataRequest, ShotDataRequestParams
 from disruption_py.utils.mappings.tokamak import Tokamak
 from disruption_py.utils.math_utils import gaussian_fit, interp1, smooth
-from disruption_py.utils.utils import without_duplicates
+from disruption_py.utils.utils import safe_cast, without_duplicates
 from disruption_py.shots.helpers.method_caching import cached_method, parameter_cached_method
 try:
     from MDSplus import mdsExceptions
@@ -1434,7 +1434,7 @@ def efit_rz2psi(params : ShotDataRequestParams, r, z, t, tree='analysis'):
         r = r.flatten()
         z = z.flatten()
         psi = np.full((len(r), len(t)), np.nan)
-        z = z.astype('float32')  # TODO: Ask if this change is necessary
+        z = safe_cast(z, 'float32')  # TODO: Ask if this change is necessary
         psirz, rgrid, zgrid, times = params.mds_conn.get_data_with_dims(r'\efit_geqdsk:psirz', tree_name=tree, dim_nums=[0, 1, 2])
         rgrid, zgrid = np.meshgrid(rgrid, zgrid) #, indexing='ij')
 

disruption_py/shots/parameter_methods/d3d/basic_parameter_methods.py

@@ -988,8 +988,8 @@ def _get_ne_te(params : ShotDataRequestParams, data_source="blessed", ts_systems
             # Place NaNs for broken channels
             lasers[laser]['te'][lasers[laser]['te'] == 0] = np.nan
             lasers[laser]['ne'][np.where(lasers[laser]['ne'] == 0)] = np.nan
-            params.logger.debug("_get_ne_te: Core bins", lasers['core']['te'].shape)
-            params.logger.debug("_get_ne_te: Tangential bins", lasers['tangential']['te'].shape)
+            params.logger.debug("_get_ne_te: Core bins {}".format(lasers['core']['te'].shape))
+            params.logger.debug("_get_ne_te: Tangential bins {}".format(lasers['tangential']['te'].shape))
         # If both systems/lasers available, combine them and interpolate the data
         # from the tangential system onto the finer (core) timebase
         if 'tangential' in lasers and lasers['tangential'] is not None: