diff --git a/astronet/tests/conftest.py b/astronet/tests/conftest.py index c1d2d88e..0bb0d1db 100644 --- a/astronet/tests/conftest.py +++ b/astronet/tests/conftest.py @@ -1,11 +1,18 @@ +import inspect +import json import subprocess import numpy as np +import pandas as pd import pytest import tensorflow as tf +from filelock import FileLock from astronet.constants import ASTRONET_WORKING_DIRECTORY as asnwd from astronet.constants import LOCAL_DEBUG +from astronet.utils import astronet_logger + +log = astronet_logger(__file__) ISA = subprocess.run( "uname -m", @@ -20,8 +27,59 @@ BATCH_SIZE = 64 -@pytest.fixture -def fixt_UGRIZY_wZ_numpy(scope="session"): +class NumpyEncoder(json.JSONEncoder): + """Special json encoder for numpy types""" + + def default(self, obj): + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.ndarray): + return obj.tolist() + return json.JSONEncoder.default(self, obj) + + +def pandas_encoder(obj): + # TODO: Reshape required to fix ValueError: Must pass 2-d input. shape=(869864, 100, 6) + # Refs: + # - https://stackoverflow.com/a/32034565/4521950 + # - https://stackoverflow.com/a/32838859/4521950 + # - https://stackoverflow.com/a/44752209/4521950 + log.critical(f"{inspect.stack()[0].function} -- Not Fully Implemented Yet") + return pd.DataFrame(obj).to_json(orient="values") + + +@pytest.fixture(scope="session") +def get_fixt_UGRIZY_wZ(tmp_path_factory, worker_id, name="fixt_UGRIZY_wZ"): + if not worker_id: + # not executing in with multiple workers, just produce the data and let + # pytest's fixture caching do its job + return fixt_UGRIZY_wZ() + + # get the temp directory shared by all workers + root_tmp_dir = tmp_path_factory.getbasetemp().parent + + fn = root_tmp_dir / "data.json" + with FileLock(str(fn) + ".lock"): + if fn.is_file(): + data = json.loads(fn.read_text()) + X_test = np.asarray(data["X_test"]) + y_test = np.asarray(data["y_test"]) + Z_test = np.asarray(data["Z_test"]) + else: + X_test, y_test, Z_test = fixt_UGRIZY_wZ() + fn.write_text( + json.dumps( + {"X_test": X_test, "y_test": y_test, "Z_test": Z_test}, + cls=NumpyEncoder, + # default=pandas_encoder, + ) + ) + return X_test, y_test, Z_test + + +def fixt_UGRIZY_wZ(): """This fixture will only be available within the scope of TestPlots""" X_test = np.load( f"{asnwd}/data/plasticc/test_set/infer/X_test.npy", @@ -33,108 +91,4 @@ def fixt_UGRIZY_wZ_numpy(scope="session"): f"{asnwd}/data/plasticc/test_set/infer/Z_test.npy", ) - inputs = [X_test, Z_test] - - return X_test, y_test, Z_test, inputs - - -@pytest.fixture -def fixt_UGRIZY_wZ(scope="session"): - """This fixture will only be available within the scope of TestPlots""" - X_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/X_test.npy", - ) - y_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/y_test.npy", - ) - Z_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/Z_test.npy", - ) - - test_input = [X_test, Z_test] - - test_ds = ( - tf.data.Dataset.from_tensor_slices( - ({"input_1": test_input[0], "input_2": test_input[1]}, y_test) - ) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - y_test_ds = ( - tf.data.Dataset.from_tensor_slices(y_test) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - if LOCAL_DEBUG is not None: - print("LOCAL_DEBUG set, reducing dataset size...") - test_ds = test_ds.take(300) - y_test_ds = y_test_ds.take(300) - - return test_ds, y_test_ds, test_input - - -@pytest.fixture -def fixt_UGRIZY_noZ(scope="session"): - """This fixture will only be available within the scope of TestPlots""" - X_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/X_test.npy", - ) - y_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/y_test.npy", - ) - - test_input = X_test - - test_ds = ( - tf.data.Dataset.from_tensor_slices((test_input, y_test)) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - y_test_ds = ( - tf.data.Dataset.from_tensor_slices(y_test) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - if LOCAL_DEBUG is not None: - print("LOCAL_DEBUG set, reducing dataset size...") - test_ds = test_ds.take(300) - y_test_ds = y_test_ds.take(300) - - return test_ds, y_test_ds - - -@pytest.fixture -def fixt_GR_noZ(scope="session"): - """This fixture will only be available within the scope of TestPlots""" - X_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/X_test.npy", - ) - y_test = np.load( - f"{asnwd}/data/plasticc/test_set/infer/y_test.npy", - ) - - X_test = X_test[:, :, 0:3:2] - test_input = X_test - - test_ds = ( - tf.data.Dataset.from_tensor_slices((test_input, y_test)) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - y_test_ds = ( - tf.data.Dataset.from_tensor_slices(y_test) - .batch(BATCH_SIZE, drop_remainder=False) - .prefetch(tf.data.AUTOTUNE) - ) - - if LOCAL_DEBUG is not None: - print("LOCAL_DEBUG set, reducing dataset size...") - test_ds = test_ds.take(300) - y_test_ds = y_test_ds.take(300) - - return test_ds, y_test_ds + return X_test, y_test, Z_test diff --git a/astronet/tests/reg/test_inference.py b/astronet/tests/reg/test_inference.py index fb38f91f..57128096 100644 --- a/astronet/tests/reg/test_inference.py +++ b/astronet/tests/reg/test_inference.py @@ -7,7 +7,9 @@ from tensorflow import keras from astronet.constants import ASTRONET_WORKING_DIRECTORY as asnwd +from astronet.constants import LOCAL_DEBUG from astronet.metrics import WeightedLogLoss +from astronet.tests.conftest import BATCH_SIZE from astronet.tinho.lite import LiteModel from astronet.utils import astronet_logger @@ -50,13 +52,41 @@ class TestInference: ), ) def test_inference_UGRIZY_wZ( - self, architecture, dataset, model_name, fixt_UGRIZY_wZ + self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ ): # Previous models were trained using numpy data as the inputs, newer models leverage # tf.data.Dataset instead for faster inference. This is a legacy requirment. # Fix ValueError of shape mismatch. - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + X_test, y_test, Z_test = get_fixt_UGRIZY_wZ + + test_input = [X_test, Z_test] + + test_ds = ( + tf.data.Dataset.from_tensor_slices( + ({"input_1": test_input[0], "input_2": test_input[1]}, y_test) + ) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + y_test_ds = ( + tf.data.Dataset.from_tensor_slices(y_test) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + if LOCAL_DEBUG is not None: + log.info("LOCAL_DEBUG set, reducing dataset size...") + test_ds = test_ds.take(300) + y_test_ds = y_test_ds.take(300) + + worker_id = ( + os.environ.get("PYTEST_XDIST_WORKER") + if "PYTEST_CURRENT_TEST" in os.environ + else 0 + ) + log.info(f"Data loaded successfully on worker: {worker_id}") model = keras.models.load_model( f"{asnwd}/astronet/{architecture}/models/{dataset}/model-{model_name}", @@ -65,7 +95,7 @@ def test_inference_UGRIZY_wZ( ) wloss = WeightedLogLoss() - y_preds = model.predict(test_inputs) + y_preds = model.predict(test_input) y_test = np.concatenate([y for y in y_test_ds], axis=0) @@ -84,10 +114,30 @@ def test_inference_UGRIZY_wZ( ), ) def test_inference_UGRIZY_noZ( - self, architecture, dataset, model_name, fixt_UGRIZY_noZ + self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ ): - test_ds, y_test_ds = fixt_UGRIZY_noZ + X_test, y_test, _ = get_fixt_UGRIZY_wZ + + test_input = X_test + + test_ds = ( + tf.data.Dataset.from_tensor_slices((test_input, y_test)) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + y_test_ds = ( + tf.data.Dataset.from_tensor_slices(y_test) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + if LOCAL_DEBUG is not None: + print("LOCAL_DEBUG set, reducing dataset size...") + test_ds = test_ds.take(300) + y_test_ds = y_test_ds.take(300) + y_test = np.concatenate([y for y in y_test_ds], axis=0) x_test = np.concatenate([x for x, y in test_ds], axis=0) @@ -125,9 +175,32 @@ def test_inference_UGRIZY_noZ( ), ), ) - def test_inference_GR_noZ(self, architecture, dataset, model_name, fixt_GR_noZ): + def test_inference_GR_noZ( + self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ + ): + + X_test, y_test, _ = get_fixt_UGRIZY_wZ + X_test = X_test[:, :, 0:3:2] + + test_input = X_test + + test_ds = ( + tf.data.Dataset.from_tensor_slices((test_input, y_test)) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + y_test_ds = ( + tf.data.Dataset.from_tensor_slices(y_test) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + if LOCAL_DEBUG is not None: + print("LOCAL_DEBUG set, reducing dataset size...") + test_ds = test_ds.take(300) + y_test_ds = y_test_ds.take(300) - test_ds, y_test_ds = fixt_GR_noZ y_test = np.concatenate([y for y in y_test_ds], axis=0) model = keras.models.load_model( @@ -152,20 +225,41 @@ def test_inference_GR_noZ(self, architecture, dataset, model_name, fixt_GR_noZ): ( "tinho", "plasticc", - "model-GR-28341-1654269564-0.5.1.dev73+g70f85f8-LL0.836.tflite", + "model-GR-noZ-28341-1654269564-0.5.1.dev73+g70f85f8-LL0.836.tflite", ), ( "tinho-quantized", "plasticc", - "quantized-model-GR-28341-1654269564-0.5.1.dev73+g70f85f8-LL0.836.tflite", + "quantized-model-GR-noZ-28341-1654269564-0.5.1.dev73+g70f85f8-LL0.836.tflite", ), ), ) def test_inference_GR_noZ_TFLITE( - self, architecture, dataset, model_name, fixt_GR_noZ + self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ ): - test_ds, y_test_ds = fixt_GR_noZ + X_test, y_test, _ = get_fixt_UGRIZY_wZ + X_test = X_test[:, :, 0:3:2] + + test_input = X_test + + test_ds = ( + tf.data.Dataset.from_tensor_slices((test_input, y_test)) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + y_test_ds = ( + tf.data.Dataset.from_tensor_slices(y_test) + .batch(BATCH_SIZE, drop_remainder=False) + .prefetch(tf.data.AUTOTUNE) + ) + + if LOCAL_DEBUG is not None: + print("LOCAL_DEBUG set, reducing dataset size...") + test_ds = test_ds.take(300) + y_test_ds = y_test_ds.take(300) + y_test = np.concatenate([y for y in y_test_ds], axis=0) x_test = np.concatenate([x for x, y in test_ds], axis=0) @@ -183,30 +277,3 @@ def test_inference_GR_noZ_TFLITE( loss = wloss(y_test, y_preds).numpy() log.info(f"LOSS tinho-quantized: {loss:.3f}") assert loss == pytest.approx(0.834, 0.001) - - # @pytest.mark.parametrize( - # ("architecture", "dataset", "model_name"), - # ( - # ( - # "tinho", - # "plasticc", - # "UGRIZY-31367-1654360237-0.5.1.dev78+g702e399.d20220604-LL0.450", - # ), - # ), - # ) - # def test_inference_with_z_tfdata(self, architecture, dataset, model_name, fixt): - - # test_ds, y_test_ds = fixt - # y_test = np.concatenate([y for y in y_test_ds], axis=0) - - # model = keras.models.load_model( - # f"{asnwd}/astronet/{architecture}/models/{dataset}/model-{model_name}", - # custom_objects={"WeightedLogLoss": WeightedLogLoss()}, - # compile=False, - # ) - - # wloss = WeightedLogLoss() - # y_preds = model.predict(test_ds) - - # if architecture == "tinho": - # assert wloss(y_test, y_preds).numpy() == pytest.approx(0.450, 0.01) diff --git a/astronet/tests/reg/test_plots.py b/astronet/tests/reg/test_plots.py index 14b06958..8601a36e 100644 --- a/astronet/tests/reg/test_plots.py +++ b/astronet/tests/reg/test_plots.py @@ -62,7 +62,7 @@ class TestPlots: """A class with common parameters, `architecture`, `dataset` and `model_name`.""" - def compute_scores(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def compute_scores(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): results_filename = ( f"{asnwd}/astronet/{architecture}/models/{dataset}/results_with_z.json" @@ -83,7 +83,7 @@ def compute_scores(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): key=lambda ev: ev["model_evaluate_on_test_loss"], ) - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ model = keras.models.load_model( f"{asnwd}/astronet/{architecture}/models/{dataset}/model-{model_name}", @@ -107,8 +107,8 @@ def compute_scores(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): def test_params(self, architecture, dataset, model_name): print("\ntest_one", architecture, dataset, model_name) - def test_fixtures(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): - print("\ntest_one", architecture, dataset, model_name, fixt_UGRIZY_wZ) + def test_fixtures(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): + print("\ntest_one", architecture, dataset, model_name, get_fixt_UGRIZY_wZ) def test_succeeds(self, architecture, dataset, model_name): fig = plt.figure() @@ -120,9 +120,9 @@ def test_succeeds(self, architecture, dataset, model_name): @pytest.mark.mpl_image_compare( # hash_library=hashlib, ) - def test_acc_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def test_acc_history(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ ( event, @@ -130,7 +130,7 @@ def test_acc_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): class_names, model, y_preds, - ) = self.compute_scores(architecture, dataset, model_name, fixt_UGRIZY_wZ) + ) = self.compute_scores(architecture, dataset, model_name, get_fixt_UGRIZY_wZ) fig = plot_acc_history( architecture, @@ -144,9 +144,9 @@ def test_acc_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): @pytest.mark.mpl_image_compare( # hash_library=hashlib, ) - def test_loss_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def test_loss_history(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ ( event, @@ -154,7 +154,7 @@ def test_loss_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): class_names, model, y_preds, - ) = self.compute_scores(architecture, dataset, model_name, fixt_UGRIZY_wZ) + ) = self.compute_scores(architecture, dataset, model_name, get_fixt_UGRIZY_wZ) fig = plot_loss_history( architecture, @@ -168,9 +168,11 @@ def test_loss_history(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): @pytest.mark.mpl_image_compare( # hash_library=hashlib, ) - def test_confusion_matrix(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def test_confusion_matrix( + self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ + ): - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ y_test = np.concatenate([y for y in y_test_ds], axis=0) ( @@ -179,7 +181,7 @@ def test_confusion_matrix(self, architecture, dataset, model_name, fixt_UGRIZY_w class_names, model, y_preds, - ) = self.compute_scores(architecture, dataset, model_name, fixt_UGRIZY_wZ) + ) = self.compute_scores(architecture, dataset, model_name, get_fixt_UGRIZY_wZ) cmap = sns.light_palette("Navy", as_cmap=True) fig = plot_confusion_matrix( @@ -198,9 +200,9 @@ def test_confusion_matrix(self, architecture, dataset, model_name, fixt_UGRIZY_w @pytest.mark.mpl_image_compare( # hash_library=hashlib, ) - def test_multiROC(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def test_multiROC(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ y_test = np.concatenate([y for y in y_test_ds], axis=0) ( @@ -209,7 +211,7 @@ def test_multiROC(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): class_names, model, y_preds, - ) = self.compute_scores(architecture, dataset, model_name, fixt_UGRIZY_wZ) + ) = self.compute_scores(architecture, dataset, model_name, get_fixt_UGRIZY_wZ) fig = plot_multiROC( architecture, @@ -226,9 +228,9 @@ def test_multiROC(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): @pytest.mark.mpl_image_compare( # hash_library=hashlib, ) - def test_multiPR(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): + def test_multiPR(self, architecture, dataset, model_name, get_fixt_UGRIZY_wZ): - test_ds, y_test_ds, test_inputs = fixt_UGRIZY_wZ + test_ds, y_test_ds, test_inputs = get_fixt_UGRIZY_wZ y_test = np.concatenate([y for y in y_test_ds], axis=0) ( @@ -237,7 +239,7 @@ def test_multiPR(self, architecture, dataset, model_name, fixt_UGRIZY_wZ): class_names, model, y_preds, - ) = self.compute_scores(architecture, dataset, model_name, fixt_UGRIZY_wZ) + ) = self.compute_scores(architecture, dataset, model_name, get_fixt_UGRIZY_wZ) fig = plot_multiPR( architecture, diff --git a/environment.yml b/environment.yml index d5a07860..39a8fbf7 100644 --- a/environment.yml +++ b/environment.yml @@ -22,6 +22,7 @@ dependencies: - pip: - astropy>=1.1.2 + - filelock - gast==0.3.3 - imageio>=2.9.0 - jax[cpu] diff --git a/pyproject.toml b/pyproject.toml index 3548cae7..7a1d5c36 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,7 @@ addopts = [ "--mpl", "--nbmake", "--log-cli-level=INFO", - "--numprocesses=1", + "--numprocesses=auto", "-ra", # "--exitfirst", ]