Skip to content

Commit 680734b

Browse files
author
STS
committed
added regression tests
1 parent 8059cd7 commit 680734b

7 files changed

Lines changed: 466 additions & 4 deletions

File tree

‎.github/workflows/build_docs.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ jobs:
1111
runs-on: ubuntu-latest
1212
strategy:
1313
matrix:
14-
python-version: [3.9]
14+
python-version: [3.11]
1515

1616
steps:
1717
- uses: actions/checkout@v3

.github/workflows/main.yml renamed to .github/workflows/regression_test.yml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# This is a basic workflow to help you get started with Actions
22

3-
name: Continous Integration
3+
name: Regression Tests
44

55
# Controls when the workflow will run
66
on:
@@ -17,7 +17,7 @@ jobs:
1717
tests_linux:
1818
strategy:
1919
matrix:
20-
python-version: [3.9]
20+
python-version: [3.11]
2121
os: [ubuntu-latest]
2222
runs-on: ${{ matrix.os }}
2323
steps:
@@ -34,7 +34,7 @@ jobs:
3434
- name: Running tests
3535
run: |
3636
source .venv/bin/activate
37-
coverage run --source=jwave -m pytest -vs
37+
coverage run --source=jwave -m pytest -vs ./tests/regression_tests
3838
coverage xml
3939
coverage html
4040
- name: "Upload coverage to Codecov"

‎.github/workflows/unit_and_integration_test.yml‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# This is a basic workflow to help you get started with Actions
2+
3+
name: Unit and Integration Tests
4+
5+
# Controls when the workflow will run
6+
on:
7+
# Triggers the workflow on push or pull request events but only for the main branch
8+
push:
9+
branches: [ main ]
10+
pull_request:
11+
branches: [ main ]
12+
13+
# Allows you to run this workflow manually from the Actions tab
14+
workflow_dispatch:
15+
16+
jobs:
17+
tests_linux:
18+
strategy:
19+
matrix:
20+
python-version: [3.11]
21+
os: [ubuntu-latest]
22+
runs-on: ${{ matrix.os }}
23+
steps:
24+
- uses: actions/checkout@v3
25+
- uses: actions/setup-python@v4
26+
with:
27+
python-version: ${{ matrix.python-version }}
28+
- name: Making virtual environment, linting
29+
run: |
30+
make testenv
31+
- name: Downloading test data
32+
run: |
33+
make get_test_data
34+
- name: Running tests
35+
run: |
36+
source .venv/bin/activate
37+
coverage run --source=jwave -m pytest -vs --ignore=tests/regression_tests
38+
coverage xml
39+
coverage html
40+
- name: "Upload coverage to Codecov"
41+
uses: codecov/codecov-action@v3
42+
with:
43+
token: ${{ secrets.CODECOV_TOKEN }}
44+
files: ./coverage.xml
45+
name: codecov-umbrella
46+
47+
# TODO: Add tests for Windows and macOS

‎tests/check_data.ipynb‎

Lines changed: 119 additions & 0 deletions
Large diffs are not rendered by default.

‎tests/regression_data/.gitkeep‎

Whitespace-only changes.

‎tests/regression_tests/test_helmholtz.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import os
2+
from functools import partial
3+
from typing import Tuple
4+
5+
import numpy as np
6+
import pytest
7+
from jax import device_put, devices, jit
8+
from jax import numpy as jnp
9+
from matplotlib import pyplot as plt
10+
from scipy.io import loadmat, savemat
11+
12+
from jwave import FourierSeries
13+
from jwave.acoustics.time_harmonic import helmholtz_solver
14+
from jwave.geometry import Domain, Medium
15+
from jwave.utils import plot_comparison
16+
17+
RELATIVE_TOLERANCE = 1e-4
18+
DIR_PATH = os.path.dirname(os.path.realpath(__file__))
19+
20+
def _make_filename(N, dx, sound_speed, density, attenuation, omega):
21+
N = str(N).replace(" ", "_")
22+
return os.path.join(
23+
DIR_PATH,
24+
"..",
25+
"regression_data",
26+
f"helmholtz_{N}_{dx}_{sound_speed}_{density}_{attenuation}_{omega}.mat".replace(" ", "_")
27+
)
28+
29+
30+
def _index_in_middle(N, span=7):
31+
return tuple(slice(Ni//2 - span, Ni//2 + span) for Ni in N)
32+
33+
34+
def _get_sos(kind, domain):
35+
match kind:
36+
case "scalar":
37+
return 1500.
38+
case "homogeneous":
39+
return FourierSeries(np.ones(domain.N) * 1500., domain)
40+
case "heterogeneous":
41+
c = np.ones(domain.N) * 1500.
42+
c[_index_in_middle(domain.N)] = 2000.
43+
return FourierSeries(c, domain)
44+
45+
46+
def _get_density(kind, domain):
47+
match kind:
48+
case "scalar":
49+
return 1000.
50+
case "homogeneous":
51+
return FourierSeries(np.ones(domain.N) * 1000., domain)
52+
case "heterogeneous":
53+
rho = np.ones(domain.N) * 1000.
54+
rho[_index_in_middle(domain.N)] = 2000.
55+
return FourierSeries(rho, domain)
56+
57+
58+
def _get_attenuation(kind, domain):
59+
match kind:
60+
case "scalar":
61+
return 0.1
62+
case "homogeneous":
63+
return FourierSeries(np.zeros(domain.N), domain)
64+
case "heterogeneous":
65+
alpha = np.zeros(domain.N)
66+
alpha[_index_in_middle(domain.N)] = 10.
67+
return FourierSeries(alpha, domain)
68+
69+
70+
@pytest.mark.parametrize("N", [(32,32), (33,31), (32,32,32), (33,31,32)])
71+
@pytest.mark.parametrize("dx", [1e-3])
72+
@pytest.mark.parametrize("sound_speed", ["scalar", "heterogeneous"])
73+
@pytest.mark.parametrize("density", ["scalar", "heterogeneous"])
74+
@pytest.mark.parametrize("attenuation", ["scalar", "heterogeneous"])
75+
@pytest.mark.parametrize("omega", [1.5e6])
76+
def test_regression_helmholtz(
77+
N, dx, sound_speed, density, attenuation, omega, reset_regression_data=False
78+
):
79+
# Setting up simulation
80+
dx = tuple([dx] * len(N))
81+
domain = Domain(N, dx)
82+
filename = _make_filename(N, dx, sound_speed, density, attenuation, omega)
83+
84+
# Making source map (dirac at center of domain)
85+
src = jnp.zeros(N, dtype=jnp.complex64)
86+
src = src.at[tuple(Ni//2 for Ni in N)].set(1.)
87+
src = FourierSeries(src, domain)
88+
89+
# Making medium
90+
sound_speed = _get_sos(sound_speed, domain)
91+
density = _get_density(density, domain)
92+
attenuation = _get_attenuation(attenuation, domain)
93+
94+
# Move everythin to cpu
95+
cpu = devices("cpu")[0]
96+
src = device_put(src, cpu)
97+
sound_speed = device_put(sound_speed, cpu)
98+
density = device_put(density, cpu)
99+
attenuation = device_put(attenuation, cpu)
100+
101+
# Initialize medium
102+
medium = Medium(domain, sound_speed, density, attenuation, pml_size=10)
103+
104+
# Run the simulation
105+
@jit
106+
def run_simulation(medium, src):
107+
return helmholtz_solver(medium, omega, src, tol=1e-5)
108+
109+
# Get field
110+
solution_field = run_simulation(medium, src)
111+
112+
# Reset regression data if needed
113+
if reset_regression_data:
114+
sol_as_numpy = np.array(solution_field.on_grid)
115+
savemat(filename, {"solution_field": sol_as_numpy})
116+
117+
# Load regression data
118+
previous_solution = loadmat(filename)["solution_field"]
119+
120+
# Make sure the solution is the same within a certain tolerance
121+
err = jnp.abs(solution_field.on_grid - previous_solution)
122+
relErr = jnp.amax(err) / jnp.amax(jnp.abs(previous_solution))
123+
print(" Relative max error = ", 100 * relErr, "%")
124+
125+
assert relErr <RELATIVE_TOLERANCE, (
126+
"Test failed, error above maximum limit of "
127+
+ str(100 *RELATIVE_TOLERANCE)
128+
+ "%"
129+
)
130+

‎tests/regression_tests/test_helmholtz_autodiff.py‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
import os
2+
from functools import partial
3+
from typing import Tuple
4+
5+
import numpy as np
6+
import pytest
7+
from jax import device_put, devices, jit, grad, value_and_grad
8+
from jax import numpy as jnp
9+
from matplotlib import pyplot as plt
10+
from scipy.io import loadmat, savemat
11+
12+
from jwave import FourierSeries
13+
from jwave.acoustics.time_harmonic import helmholtz_solver
14+
from jwave.geometry import Domain, Medium
15+
from jwave.utils import plot_comparison
16+
17+
RELATIVE_TOLERANCE = 1e-4
18+
DIR_PATH = os.path.dirname(os.path.realpath(__file__))
19+
20+
def _make_filename(N, dx, sound_speed, density, attenuation, omega):
21+
N = str(N).replace(" ", "_")
22+
return os.path.join(
23+
DIR_PATH,
24+
"..",
25+
"regression_data",
26+
f"helmholtz_autodiff_{N}_{dx}_{sound_speed}_{density}_{attenuation}_{omega}.mat".replace(" ", "_")
27+
)
28+
29+
30+
def _index_in_middle(N, span=7):
31+
return tuple(slice(Ni//2 - span, Ni//2 + span) for Ni in N)
32+
33+
34+
def _get_sos(kind, domain):
35+
match kind:
36+
case "scalar":
37+
return 1500.
38+
case "homogeneous":
39+
return FourierSeries(np.ones(domain.N) * 1500., domain)
40+
case "heterogeneous":
41+
c = np.ones(domain.N) * 1500.
42+
c[_index_in_middle(domain.N)] = 2000.
43+
return FourierSeries(c, domain)
44+
45+
46+
def _get_density(kind, domain):
47+
match kind:
48+
case "scalar":
49+
return 1000.
50+
case "homogeneous":
51+
return FourierSeries(np.ones(domain.N) * 1000., domain)
52+
case "heterogeneous":
53+
rho = np.ones(domain.N) * 1000.
54+
rho[_index_in_middle(domain.N)] = 2000.
55+
return FourierSeries(rho, domain)
56+
57+
58+
def _get_attenuation(kind, domain):
59+
match kind:
60+
case "scalar":
61+
return 0.1
62+
case "homogeneous":
63+
return FourierSeries(np.zeros(domain.N), domain)
64+
case "heterogeneous":
65+
alpha = np.zeros(domain.N)
66+
alpha[_index_in_middle(domain.N)] = 10.
67+
return FourierSeries(alpha, domain)
68+
69+
70+
@pytest.mark.parametrize("N", [(48,48), (49,47), (32,32,32), (33,31,32)])
71+
@pytest.mark.parametrize("dx", [1e-3])
72+
@pytest.mark.parametrize("sound_speed", ["heterogeneous"])
73+
@pytest.mark.parametrize("density", ["heterogeneous"])
74+
@pytest.mark.parametrize("attenuation", ["heterogeneous"])
75+
@pytest.mark.parametrize("omega", [1e6])
76+
def test_regression_helmholtz(
77+
N, dx, sound_speed, density, attenuation, omega, reset_regression_data=False
78+
):
79+
# Setting up simulation
80+
dx = tuple([dx] * len(N))
81+
domain = Domain(N, dx)
82+
filename = _make_filename(N, dx, sound_speed, density, attenuation, omega)
83+
84+
# Making source map (dirac at center of domain)
85+
src = jnp.zeros(N, dtype=jnp.complex64)
86+
src = src.at[tuple(11 for Ni in N)].set(1.)
87+
src = FourierSeries(src, domain)
88+
89+
# Making medium
90+
sound_speed = _get_sos(sound_speed, domain)
91+
density = _get_density(density, domain)
92+
attenuation = _get_attenuation(attenuation, domain)
93+
94+
# Move everythin to cpu
95+
cpu = devices("cpu")[0]
96+
src = device_put(src, cpu)
97+
sound_speed = device_put(sound_speed, cpu)
98+
density = device_put(density, cpu)
99+
attenuation = device_put(attenuation, cpu)
100+
101+
@jit
102+
@partial(grad, argnums=[0,1,2,3,4], has_aux=True)
103+
def loss_fn(
104+
sound_speed: FourierSeries,
105+
density: FourierSeries,
106+
attenuation: FourierSeries,
107+
omega: float,
108+
src: FourierSeries,
109+
):
110+
# This tries to maximize the amplitude of the field at a given point
111+
medium = Medium(src.domain, sound_speed, density, attenuation, pml_size=10)
112+
solution_field = helmholtz_solver(medium, omega, src, tol=1e-5).on_grid
113+
max_point = [-11] * len(src.domain.N)
114+
max_point = tuple(max_point)
115+
return -jnp.sum(jnp.abs(solution_field[max_point])), solution_field
116+
117+
# Get gradients
118+
gradients, field = loss_fn(
119+
sound_speed,
120+
density,
121+
attenuation,
122+
omega,
123+
src
124+
)
125+
126+
# Make them numpy arrays
127+
sos_gradient = np.array(gradients[0].on_grid)
128+
density_gradient = np.array(gradients[1].on_grid)
129+
attenuation_gradient = np.array(gradients[2].on_grid)
130+
omega_gradient = np.array(gradients[3])
131+
src_gradient = np.array(gradients[4].on_grid)
132+
133+
# Reset regression data if needed
134+
if reset_regression_data:
135+
field = np.array(field)
136+
savemat(filename, {
137+
"sos_gradient": sos_gradient,
138+
"density_gradient": density_gradient,
139+
"attenuation_gradient": attenuation_gradient,
140+
"omega_gradient": omega_gradient,
141+
"src_gradient": src_gradient,
142+
"field": field
143+
})
144+
145+
# Load regression data
146+
matfile = loadmat(filename)
147+
148+
# Check each one of them
149+
err_fun = lambda x,y: jnp.amax(jnp.abs(x-y))/ jnp.amax(jnp.abs(y))
150+
max_rel_error = max([
151+
err_fun(sos_gradient, matfile["sos_gradient"]),
152+
err_fun(density_gradient, matfile["density_gradient"]),
153+
err_fun(attenuation_gradient, matfile["attenuation_gradient"]),
154+
err_fun(omega_gradient, matfile["omega_gradient"]),
155+
err_fun(src_gradient, matfile["src_gradient"]),
156+
])
157+
158+
# Make sure the solution is the same within a certain tolerance
159+
print(" Relative max error = ", 100 * max_rel_error, "%")
160+
161+
assert max_rel_error <RELATIVE_TOLERANCE, (
162+
"Test failed, error above maximum limit of "
163+
+ str(100 *RELATIVE_TOLERANCE)
164+
+ "%"
165+
)
166+

0 commit comments

Comments
 (0)