Unit tests for Crystal

.gitignore

@@ -3,8 +3,11 @@ __pycache__
 custom-poling/__pycache__/
 custom_poling/__pycache__/__init__.cpython-38.pyc
 custom_poling/__pycache__/__init__.cpython-39.pyc
+custom_poling/core/__pycache__
+custom_poling/utils/__pycache__
 .DS_Store
 .python-version
+.coverage
 
 # Distribution / packaging
 pyproject.toml

custom_poling/core/crystal.py

@@ -34,13 +34,13 @@ def compute_pmf(self, domain_configuration, k_array):
 
         Args:
             domain_configuration (list of int): elements of list must be +1 or -1
-            k_array (array of floats)
+            k_array (array of floats): k values for which to compute the PMF 
 
         Returns:
             PMF as an array of floats
         """
         self.domain_configuration = domain_configuration
-        self.k_array = k_array
+        self.k_array = k_array 
         crystal_pmf = pmf(self.domain_walls, self.domain_configuration, self.k_array)
         return crystal_pmf
 

custom_poling/utils/pmf.py

@@ -1,7 +1,10 @@
 import numpy as np
 
-def pmf(domain_walls, domain_configuration, k):
-    pmf_one_domain = lambda z1,z2:2*np.pi*1j*(np.exp(1j*k*z1)-np.exp(1j*k*z2))/(k*2*np.pi)
+def pmf(domain_walls, domain_configuration, k_array):
+    """Returns the phasematching function (PMF) as a function of k for a given domain_configuration."""
+    if np.any(k_array==0.):
+        raise ValueError('Sorry, no numbers equal to zero inside k_array')
+    pmf_one_domain = lambda z1,z2:2*np.pi*1j*(np.exp(1j*k_array*z1)-np.exp(1j*k_array*z2))/(k_array*2*np.pi)
     pmf = 0
     for idx in range(len(domain_configuration)):
         pmf = pmf + domain_configuration[idx]*pmf_one_domain(domain_walls[idx],domain_walls[idx+1])

tests/test_crystal.py

@@ -0,0 +1,69 @@
+import numpy as np
+import os
+import pytest
+
+from custom_poling.core.custom_crystal import Crystal
+
+def test_crystal_attributes() -> None:
+    domain_width = 1
+    number_domains = 10
+    z0 = 0
+    crystal = Crystal(domain_width, number_domains, z0)
+    assert crystal.domain_width == domain_width
+    assert crystal.number_domains == number_domains
+    assert crystal.z0 == z0
+    assert crystal.length == domain_width * number_domains
+    assert np.array_equal(crystal.domain_walls, np.arange(z0, z0 + (number_domains + 1) * domain_width, domain_width))
+    assert np.array_equal(crystal.domain_middles, (crystal.domain_walls + domain_width/2)[0:-1])
+
+def test_crystal_compute_pmf_array_length() -> None:
+    domain_width = 1
+    number_domains = 10
+    z0 = 0
+    crystal = Crystal(domain_width, number_domains, z0)
+    domain_configuration = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1]
+    k_array = np.linspace(-1, 1, 100)
+    crystal_pmf = crystal.compute_pmf(domain_configuration, k_array)
+    assert len(crystal_pmf) == len(k_array)
+
+def test_crystal_compute_pmf_array_values() -> None:
+    domain_width = 1
+    number_domains = 10
+    z0 = 0
+    crystal = Crystal(domain_width, number_domains, z0)
+    domain_configuration = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1]
+    k_array = np.linspace(-1, 1, 6)
+    crystal_pmf = crystal.compute_pmf(domain_configuration, k_array)
+    expected_crystal_pmf=np.array([1.00468936-0.29720009j, 0.02053462-0.14405557j,
+       0.71044314+0.4561703j , 0.71044314-0.4561703j ,
+       0.02053462+0.14405557j, 1.00468936+0.29720009j])
+    assert np.allclose(crystal_pmf,expected_crystal_pmf,rtol=1e-07)
+
+def test_crystal_compute_pmf_invalid_karray() -> None:
+    with pytest.raises(ValueError):
+        domain_width = 1
+        number_domains = 10
+        z0 = 0
+        crystal = Crystal(domain_width, number_domains, z0)
+        domain_configuration = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1]
+        k_array = np.linspace(-1, 1, 5)
+        crystal.compute_pmf(domain_configuration, k_array)
+
+def test_crystal_plot_domains_and_save() -> None:
+    domain_width = 1
+    number_domains = 10
+    z0 = 0
+    crystal = Crystal(domain_width, number_domains, z0)
+    domain_configuration = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1]
+    crystal.plot_domains(domain_configuration, n_max=5, show=False, save_as='test_crystal.png', fix_ticks=True)
+    assert os.path.isfile('test_crystal.png')
+    os.remove('test_crystal.png')
+
+# def test_crystal_plot_domains_and_show() -> None:
+#     domain_width = 1
+#     number_domains = 10
+#     z0 = 0
+#     crystal = Crystal(domain_width, number_domains, z0)
+#     domain_configuration = [1, -1, 1, -1, 1, -1, 1, -1, 1, -1]
+#     crystal.plot_domains(domain_configuration, n_max=5, show=True, save_as='test_crystal.png', fix_ticks=True)
+#     # TBD