from pathlib import Path from typing import Any from unittest import TestCase import numpy as np import shutil import unittest import tempfile import subprocess import itertools import argparse class GrayTest: inputs: Path = None # directory of the input files reference: Path = None # directory of the reference outputs candidate: Path = None # directory of the candidate outputs # Extra parameters to pass to gray gray_params: dict[str, Any] = {} @classmethod def setUpClass(cls): ''' Sets up the test case ''' # directory of the test case base = Path().joinpath(*cls.__module__.split('.')) if cls.inputs is None: cls.inputs = base / 'inputs' if cls.reference is None: cls.reference = base / 'outputs' # temporary directory holding the candidate outputs cls._tempdir = tempfile.mkdtemp(prefix=f'gray-test-{base.name}.') cls.candidate = Path(cls._tempdir) # replace reference with candidate if options.update: print() print('Setting new reference for ' + cls.__module__) cls.candidate = cls.reference # run gray to generate the candidate outputs proc = run_gray(cls.inputs, cls.candidate, params=cls.gray_params, binary=options.binary) # 0: all good, 1: input errors, >1: simulation errors assert proc.returncode != 1, 'gray failed with exit code 1' # store the stderr for manual inspection with open(str(cls.candidate / 'log'), 'w') as log: log.write(proc.stderr) @classmethod def tearDownClass(cls): ''' Clean up after all tests ''' # remove temporary directory if cls._passed or not options.keep_failed: shutil.rmtree(cls._tempdir) else: print() print('Some tests failed: preserving outputs in', cls._tempdir) def run(self, result: unittest.runner.TextTestResult): ''' Override to store the test results for tearDownClass ''' TestCase.run(self, result) self.__class__._passed = result.failures == [] def test_eccd_values(self): ''' Comparing the ECCD values ''' from collections import defaultdict try: ref = load_table(self.reference / 'summary.7.txt') cand = load_table(self.candidate / 'summary.7.txt') except FileNotFoundError: raise unittest.SkipTest("ECCD results not available") # precision as number of decimal places prec = defaultdict(lambda: 3, [ ('dPdV_peak', -2), ('dPdV_max', -2), ('J_φ_peak', -2), ('J_φ_max', -2), ('s_max', -1), ('χ', -1), ('ψ', -1), ]) for val in ref.dtype.names: with self.subTest(value=val): for i, ray in enumerate(ref['index_rt']): self.assertAlmostEqual( ref[val][i], cand[val][i], prec[val], msg=f"{val} changed (ray {int(ray)})") def test_final_position(self): ''' Comparing the final position of the central ray ''' ref = load_table(self.reference / 'central-ray.4.txt') cand = load_table(self.candidate / 'central-ray.4.txt') # coordinates self.assertAlmostEqual(ref['R'][-1], cand['R'][-1], 1) self.assertAlmostEqual(ref['z'][-1], cand['z'][-1], 1) self.assertAlmostEqual(ref['φ'][-1], cand['φ'][-1], 2) # optical path length self.assertAlmostEqual(ref['s'][-1], cand['s'][-1], 1) def test_final_direction(self): ''' Comparing the final direction of the central ray ''' ref = load_table(self.reference / 'central-ray.4.txt') cand = load_table(self.candidate / 'central-ray.4.txt') self.assertAlmostEqual(ref['N_⊥'][-1], cand['N_⊥'][-1], 1) self.assertAlmostEqual(ref['N_∥'][-1], cand['N_∥'][-1], 1) def test_beam_shape(self): ''' Comparing the final beam shape ''' try: ref = load_table(self.reference / 'beam-shape-final.9.txt') cand = load_table(self.candidate / 'beam-shape-final.9.txt') except FileNotFoundError: raise unittest.SkipTest("Beam shape info not available") if options.visual: import matplotlib.pyplot as plt plt.subplot(aspect='equal') plt.title(self.__module__ + '.test_beam_shape') plt.xlabel('$x$ / cm') plt.ylabel('$y$ / cm') plt.scatter(ref['x'], ref['y'], c='red', marker='_', label='reference') plt.scatter(cand['x'], cand['y'], c='green', alpha=0.6, marker='+', label='candidate') plt.legend() plt.show() for ref, cand in zip(ref, cand): with self.subTest(ray=(int(ref['j']), int(ref['k']))): self.assertAlmostEqual(ref['x'], cand['x'], 1) self.assertAlmostEqual(ref['y'], cand['y'], 1) def test_error_biased(self): ''' Test for a proportionality between Λ and any of X, Y, N∥ ''' data = load_table(self.candidate / 'central-ray.4.txt') # restrict to within the plasma, half of the first pass in_plasma = data['X'] > 0 first_pass = data['index_rt'] == data['index_rt'].min() data = data[in_plasma & first_pass] data = data[:int(data.size // 2)] if data.size < 2: self.skipTest("There is no plasma") if options.visual: import matplotlib.pyplot as plt left = plt.subplot() plt.title(self.__module__ + '.test_error_biased') left.set_xlabel('$s$ / cm') left.set_ylabel('$Λ$', color='xkcd:ocean blue') left.tick_params(axis='y', labelcolor='xkcd:ocean blue') left.plot(data['s'], data['Λ_r'], color='xkcd:ocean blue') right1 = left.twinx() right1.set_ylabel('$X$', color='xkcd:orange') right1.tick_params(axis='y', labelcolor='xkcd:orange') right1.plot(data['s'], data['X'], color='xkcd:orange') right2 = left.twinx() right2.set_ylabel('$Y$', color='xkcd:vermillion') right2.tick_params(axis='y', labelcolor='xkcd:vermillion') right2.plot(data['s'], data['Y'], color='xkcd:vermillion') right2.spines["right"].set_position(("axes", 1.1)) right3 = left.twinx() right3.set_ylabel('$N_∥$', color='xkcd:green') right3.tick_params(axis='y', labelcolor='xkcd:green') right3.spines["right"].set_position(("axes", 1.2)) right3.plot(data['s'], data['N_∥'], color='xkcd:green') plt.subplots_adjust(right=0.78) plt.show() err = data['Λ_r'].var() / 10 self.assertGreater(err, 0, msg="Λ is exactly constant") for var in ['X', 'Y', 'N_⊥']: # Minimise the χ²(k) = |(Λ_r - k⋅var) / err|² / (n - 1) # The solution is simply: k = (Λ⋅var)/var⋅var k = np.dot(data['Λ_r'], data[var]) / np.linalg.norm(data[var])**2 with self.subTest(var=var): res = (data['Λ_r'] - k*data[var]) / err χ2 = np.linalg.norm(res)**2 / (data.size - 1) self.assertGreater(χ2, 1) # Command line options options = argparse.Namespace() def get_basedir(module: str) -> Path: """ Given a module name (es. tests.03-TCV) returns its base directory as a path (es. tests/03-TCV). """ return Path().joinpath(*module.split('.')) def run_gray(inputs: Path, outputs: Path, # extra gray parameters params: dict[str, Any] = {}, # which tables to generate tables: list[int] = [4, 7, 8, 9, 48, 33, 70, 71], # which gray binary to use binary: str = 'gray', # extra options options: [str] = [] ) -> subprocess.CompletedProcess: ''' Runs gray on the inputs from the `inputs` directory and storing the results in the `outputs` directory. ''' outputs.mkdir(exist_ok=True, parents=True) params = [['-g', f'{k}={v}'] for k, v in params.items()] args = [ binary, '-c', str(inputs / 'gray.ini'), '-t', ','.join(map(str, tables)), '-o', str(outputs), '-v' ] + list(itertools.chain(*params)) + options proc = subprocess.run(args, capture_output=True, text=True) print() if proc.returncode != 0: # show the log on errors print(f'Errors occurred (exit status {proc.returncode}), showing log:') print(*proc.args) print(proc.stderr) print(proc.stdout) return proc def load_table(fname: Path) -> np.array: ''' Loads a GRAY output file as a structured numpy array (columns are named as in the file header) ''' return np.genfromtxt(fname, names=True, skip_header=21, ndmin=1)