Skip to content

Commit

Permalink
refactor: Added type hints and docstrings to helper functions. Remove…
Browse files Browse the repository at this point in the history
…d unused helper functions
  • Loading branch information
pabloitu committed Aug 23, 2024
1 parent 94bf14c commit 5d2011c
Showing 1 changed file with 67 additions and 67 deletions.
134 changes: 67 additions & 67 deletions tests/test_plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,73 +223,73 @@ def savefig(self, ax, name):
# gc.collect()
#
#
# class TestPlotMagnitudeHistogram(TestPlots):
#
# def setUp(self):
#
# def gr_dist(num_events, mag_min=3.0, mag_max=8.0, b_val=1.0):
# U = numpy.random.uniform(0, 1, num_events)
# magnitudes = mag_min - (1.0 / b_val) * numpy.log10(1 - U)
# magnitudes = magnitudes[magnitudes <= mag_max]
# return magnitudes
#
# self.mock_forecast = [MagicMock(), MagicMock(), MagicMock()]
# for i in self.mock_forecast:
# i.get_magnitudes.return_value = gr_dist(5000)
#
# self.mock_cat = MagicMock()
# self.mock_cat.get_magnitudes.return_value = gr_dist(500, b_val=1.2)
# self.mock_cat.get_number_of_events.return_value = 500
# self.mock_cat.region.magnitudes = numpy.arange(3.0, 8.0, 0.1)
# self.save_dir = os.path.join(os.path.dirname(__file__), "artifacts", "plots")
#
# cat_file_m5 = os.path.join(
# self.artifacts,
# "example_csep2_forecasts",
# "Catalog",
# "catalog.json",
# )
# self.comcat = catalogs.CSEPCatalog.load_json(cat_file_m5)
# forecast_file = os.path.join(
# self.artifacts,
# "example_csep2_forecasts",
# "Forecasts",
# "ucerf3-landers_short.csv",
# )
#
# self.stochastic_event_sets = csep.load_catalog_forecast(forecast_file)
#
# os.makedirs(self.save_dir, exist_ok=True)
#
# def test_plot_magnitude_histogram_basic(self):
# # Test with basic arguments
# ax = plot_magnitude_histogram(self.mock_forecast,
# self.mock_cat, show=show_plots,
# density=True)
#
# # Verify that magnitudes were retrieved
# for catalog in self.mock_forecast:
# catalog.get_magnitudes.assert_called_once()
# self.mock_cat.get_magnitudes.assert_called_once()
# self.mock_cat.get_number_of_events.assert_called_once()
# ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram.png"))
#
# def test_plot_magnitude_histogram_ucerf(self):
# # Test with basic arguments
# ax = plot_magnitude_histogram(self.stochastic_event_sets, self.comcat,
# show=show_plots)
#
# # # Verify that magnitudes were retrieved
# # for catalog in self.stochastic_event_sets:
# # catalog.get_magnitudes.assert_called_once()
# # self.comcat.get_magnitudes.assert_called_once()
# # self.comcat.get_number_of_events.assert_called_once()
# ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram_ucerf.png"))
#
# def tearDown(self):
# plt.close("all")
# gc.collect()
#
class TestPlotMagnitudeHistogram(TestPlots):

def setUp(self):

def gr_dist(num_events, mag_min=3.0, mag_max=8.0, b_val=1.0):
U = numpy.random.uniform(0, 1, num_events)
magnitudes = mag_min - (1.0 / b_val) * numpy.log10(1 - U)
magnitudes = magnitudes[magnitudes <= mag_max]
return magnitudes

self.mock_forecast = [MagicMock(), MagicMock(), MagicMock()]
for i in self.mock_forecast:
i.get_magnitudes.return_value = gr_dist(5000)

self.mock_cat = MagicMock()
self.mock_cat.get_magnitudes.return_value = gr_dist(500, b_val=1.2)
self.mock_cat.get_number_of_events.return_value = 500
self.mock_cat.region.magnitudes = numpy.arange(3.0, 8.0, 0.1)
self.save_dir = os.path.join(os.path.dirname(__file__), "artifacts", "plots")

cat_file_m5 = os.path.join(
self.artifacts,
"example_csep2_forecasts",
"Catalog",
"catalog.json",
)
self.comcat = catalogs.CSEPCatalog.load_json(cat_file_m5)
forecast_file = os.path.join(
self.artifacts,
"example_csep2_forecasts",
"Forecasts",
"ucerf3-landers_short.csv",
)

self.stochastic_event_sets = csep.load_catalog_forecast(forecast_file)

os.makedirs(self.save_dir, exist_ok=True)

def test_plot_magnitude_histogram_basic(self):
# Test with basic arguments
ax = plot_magnitude_histogram(self.mock_forecast,
self.mock_cat, show=show_plots,
density=True)

# Verify that magnitudes were retrieved
for catalog in self.mock_forecast:
catalog.get_magnitudes.assert_called_once()
self.mock_cat.get_magnitudes.assert_called_once()
self.mock_cat.get_number_of_events.assert_called_once()
ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram.png"))

def test_plot_magnitude_histogram_ucerf(self):
# Test with basic arguments
ax = plot_magnitude_histogram(self.stochastic_event_sets, self.comcat,
show=show_plots)

# # Verify that magnitudes were retrieved
# for catalog in self.stochastic_event_sets:
# catalog.get_magnitudes.assert_called_once()
# self.comcat.get_magnitudes.assert_called_once()
# self.comcat.get_number_of_events.assert_called_once()
ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram_ucerf.png"))

def tearDown(self):
plt.close("all")
gc.collect()

#
class TestPlotDistributionTests(TestPlots):

Expand Down

0 comments on commit 5d2011c

Please sign in to comment.