From 2e17704fdc18937b38c2dc33a55120285b7cd776 Mon Sep 17 00:00:00 2001 From: Abhimanyu Susobhanan Date: Mon, 9 Oct 2023 16:44:56 -0500 Subject: [PATCH] fix wavex_setup and dmwavex_setup --- src/pint/utils.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/pint/utils.py b/src/pint/utils.py index c4410618a..9fcdb31dc 100644 --- a/src/pint/utils.py +++ b/src/pint/utils.py @@ -1328,8 +1328,9 @@ def wavex_setup(model, T_span, freqs=None, n_freqs=None): "Both freqs and n_freqs are specified. Only one or the other should be used" ) - if n_freqs <= 0: + if n_freqs is not None and n_freqs <= 0: raise ValueError("Must use a non-zero number of wave frequencies") + model.add_component(WaveX()) if isinstance(T_span, u.quantity.Quantity): T_span.to(u.d) @@ -1356,11 +1357,11 @@ def wavex_setup(model, T_span, freqs=None, n_freqs=None): if n_freqs is not None: if n_freqs == 1: - wave_freq = 2.0 * np.pi / T_span + wave_freq = 1 / T_span model.WXFREQ_0001.quantity = wave_freq else: wave_numbers = np.arange(1, n_freqs + 1) - wave_freqs = 2.0 * np.pi * wave_numbers / T_span + wave_freqs = wave_numbers / T_span model.WXFREQ_0001.quantity = wave_freqs[0] model.components["WaveX"].add_wavex_components(wave_freqs[1:]) return model.components["WaveX"].get_indices() @@ -1408,8 +1409,9 @@ def dmwavex_setup(model, T_span, freqs=None, n_freqs=None): "Both freqs and n_freqs are specified. Only one or the other should be used" ) - if n_freqs <= 0: + if n_freqs is not None and n_freqs <= 0: raise ValueError("Must use a non-zero number of wave frequencies") + model.add_component(DMWaveX()) if isinstance(T_span, u.quantity.Quantity): T_span.to(u.d) @@ -1436,11 +1438,11 @@ def dmwavex_setup(model, T_span, freqs=None, n_freqs=None): if n_freqs is not None: if n_freqs == 1: - wave_freq = 2.0 * np.pi / T_span + wave_freq = 1 / T_span model.DMWXFREQ_0001.quantity = wave_freq else: wave_numbers = np.arange(1, n_freqs + 1) - wave_freqs = 2.0 * np.pi * wave_numbers / T_span + wave_freqs = wave_numbers / T_span model.DMWXFREQ_0001.quantity = wave_freqs[0] model.components["DMWaveX"].add_dmwavex_components(wave_freqs[1:]) return model.components["DMWaveX"].get_indices()