Skip to content

Commit 5d2011c

Browse files
committed
refactor: Added type hints and docstrings to helper functions. Removed unused helper functions
1 parent 94bf14c commit 5d2011c

File tree

1 file changed

+67
-67
lines changed

1 file changed

+67
-67
lines changed

tests/test_plots.py

+67-67
Original file line numberDiff line numberDiff line change
@@ -223,73 +223,73 @@ def savefig(self, ax, name):
223223
# gc.collect()
224224
#
225225
#
226-
# class TestPlotMagnitudeHistogram(TestPlots):
227-
#
228-
# def setUp(self):
229-
#
230-
# def gr_dist(num_events, mag_min=3.0, mag_max=8.0, b_val=1.0):
231-
# U = numpy.random.uniform(0, 1, num_events)
232-
# magnitudes = mag_min - (1.0 / b_val) * numpy.log10(1 - U)
233-
# magnitudes = magnitudes[magnitudes <= mag_max]
234-
# return magnitudes
235-
#
236-
# self.mock_forecast = [MagicMock(), MagicMock(), MagicMock()]
237-
# for i in self.mock_forecast:
238-
# i.get_magnitudes.return_value = gr_dist(5000)
239-
#
240-
# self.mock_cat = MagicMock()
241-
# self.mock_cat.get_magnitudes.return_value = gr_dist(500, b_val=1.2)
242-
# self.mock_cat.get_number_of_events.return_value = 500
243-
# self.mock_cat.region.magnitudes = numpy.arange(3.0, 8.0, 0.1)
244-
# self.save_dir = os.path.join(os.path.dirname(__file__), "artifacts", "plots")
245-
#
246-
# cat_file_m5 = os.path.join(
247-
# self.artifacts,
248-
# "example_csep2_forecasts",
249-
# "Catalog",
250-
# "catalog.json",
251-
# )
252-
# self.comcat = catalogs.CSEPCatalog.load_json(cat_file_m5)
253-
# forecast_file = os.path.join(
254-
# self.artifacts,
255-
# "example_csep2_forecasts",
256-
# "Forecasts",
257-
# "ucerf3-landers_short.csv",
258-
# )
259-
#
260-
# self.stochastic_event_sets = csep.load_catalog_forecast(forecast_file)
261-
#
262-
# os.makedirs(self.save_dir, exist_ok=True)
263-
#
264-
# def test_plot_magnitude_histogram_basic(self):
265-
# # Test with basic arguments
266-
# ax = plot_magnitude_histogram(self.mock_forecast,
267-
# self.mock_cat, show=show_plots,
268-
# density=True)
269-
#
270-
# # Verify that magnitudes were retrieved
271-
# for catalog in self.mock_forecast:
272-
# catalog.get_magnitudes.assert_called_once()
273-
# self.mock_cat.get_magnitudes.assert_called_once()
274-
# self.mock_cat.get_number_of_events.assert_called_once()
275-
# ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram.png"))
276-
#
277-
# def test_plot_magnitude_histogram_ucerf(self):
278-
# # Test with basic arguments
279-
# ax = plot_magnitude_histogram(self.stochastic_event_sets, self.comcat,
280-
# show=show_plots)
281-
#
282-
# # # Verify that magnitudes were retrieved
283-
# # for catalog in self.stochastic_event_sets:
284-
# # catalog.get_magnitudes.assert_called_once()
285-
# # self.comcat.get_magnitudes.assert_called_once()
286-
# # self.comcat.get_number_of_events.assert_called_once()
287-
# ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram_ucerf.png"))
288-
#
289-
# def tearDown(self):
290-
# plt.close("all")
291-
# gc.collect()
292-
#
226+
class TestPlotMagnitudeHistogram(TestPlots):
227+
228+
def setUp(self):
229+
230+
def gr_dist(num_events, mag_min=3.0, mag_max=8.0, b_val=1.0):
231+
U = numpy.random.uniform(0, 1, num_events)
232+
magnitudes = mag_min - (1.0 / b_val) * numpy.log10(1 - U)
233+
magnitudes = magnitudes[magnitudes <= mag_max]
234+
return magnitudes
235+
236+
self.mock_forecast = [MagicMock(), MagicMock(), MagicMock()]
237+
for i in self.mock_forecast:
238+
i.get_magnitudes.return_value = gr_dist(5000)
239+
240+
self.mock_cat = MagicMock()
241+
self.mock_cat.get_magnitudes.return_value = gr_dist(500, b_val=1.2)
242+
self.mock_cat.get_number_of_events.return_value = 500
243+
self.mock_cat.region.magnitudes = numpy.arange(3.0, 8.0, 0.1)
244+
self.save_dir = os.path.join(os.path.dirname(__file__), "artifacts", "plots")
245+
246+
cat_file_m5 = os.path.join(
247+
self.artifacts,
248+
"example_csep2_forecasts",
249+
"Catalog",
250+
"catalog.json",
251+
)
252+
self.comcat = catalogs.CSEPCatalog.load_json(cat_file_m5)
253+
forecast_file = os.path.join(
254+
self.artifacts,
255+
"example_csep2_forecasts",
256+
"Forecasts",
257+
"ucerf3-landers_short.csv",
258+
)
259+
260+
self.stochastic_event_sets = csep.load_catalog_forecast(forecast_file)
261+
262+
os.makedirs(self.save_dir, exist_ok=True)
263+
264+
def test_plot_magnitude_histogram_basic(self):
265+
# Test with basic arguments
266+
ax = plot_magnitude_histogram(self.mock_forecast,
267+
self.mock_cat, show=show_plots,
268+
density=True)
269+
270+
# Verify that magnitudes were retrieved
271+
for catalog in self.mock_forecast:
272+
catalog.get_magnitudes.assert_called_once()
273+
self.mock_cat.get_magnitudes.assert_called_once()
274+
self.mock_cat.get_number_of_events.assert_called_once()
275+
ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram.png"))
276+
277+
def test_plot_magnitude_histogram_ucerf(self):
278+
# Test with basic arguments
279+
ax = plot_magnitude_histogram(self.stochastic_event_sets, self.comcat,
280+
show=show_plots)
281+
282+
# # Verify that magnitudes were retrieved
283+
# for catalog in self.stochastic_event_sets:
284+
# catalog.get_magnitudes.assert_called_once()
285+
# self.comcat.get_magnitudes.assert_called_once()
286+
# self.comcat.get_number_of_events.assert_called_once()
287+
ax.figure.savefig(os.path.join(self.save_dir, "magnitude_histogram_ucerf.png"))
288+
289+
def tearDown(self):
290+
plt.close("all")
291+
gc.collect()
292+
293293
#
294294
class TestPlotDistributionTests(TestPlots):
295295

0 commit comments

Comments
 (0)