Merge pull request #21 from qpv-research-group/new_tests

New tests

Author: phoebe-p

rayflare/matrix_formalism/ideal_cases.py

@@ -23,47 +23,45 @@ def lambertian_matrix(angle_vector, theta_intv, surf_name, structpath,
         savepath_RT = os.path.join(structpath, surf_name + front_or_rear + 'RT.npz')
         savepath_A = os.path.join(structpath, surf_name + front_or_rear + 'A.npz')
 
-    if os.path.isfile(savepath_RT) and save:
-        print('Existing angular redistribution matrices found')
-        allArray = load_npz(savepath_RT)
+        if os.path.isfile(savepath_RT):
+            print('Existing angular redistribution matrices found')
+            allArray = load_npz(savepath_RT)
+            return allArray
 
-    else:
-
-        theta_values = np.unique(angle_vector[angle_vector[:,1] < np.pi/2,1])
-        dtheta = np.diff(theta_intv[theta_intv <= np.pi/2])
 
-        dP = np.cos(theta_values)*dtheta
+    theta_values = np.unique(angle_vector[angle_vector[:,1] < np.pi/2,1])
+    dtheta = np.diff(theta_intv[theta_intv <= np.pi/2])
 
-        # matrix has indexing (out, in): row picks out 'out' entry, column picks out which v0 element
+    dP = np.cos(theta_values)*dtheta
 
-        # since it doesn't matter what the incidence angle is for Lambertian scattering, all the columns rows be identical!
+    # matrix has indexing (out, in): row picks out 'out' entry, column picks out which v0 element
 
-        # how many phi entries are there for each theta?
+    # since it doesn't matter what the incidence angle is for Lambertian scattering, all the columns rows be identical!
 
-        n_phis = [np.sum(angle_vector[:,1] == theta) for theta in theta_values]
+    # how many phi entries are there for each theta?
 
-        column = [x for sublist in [[dP[i1]/n]*n for i1, n in enumerate(n_phis)] for x in sublist]
+    n_phis = [np.sum(angle_vector[:,1] == theta) for theta in theta_values]
 
-        whole_matrix = np.vstack([column]*int(len(angle_vector)/2)).T
+    column = [x for sublist in [[dP[i1]/n]*n for i1, n in enumerate(n_phis)] for x in sublist]
 
-        # renormalize (rounding errors)
+    whole_matrix = np.vstack([column]*int(len(angle_vector)/2)).T
 
-        whole_matrix_R = whole_matrix/np.sum(whole_matrix,0)
+    # renormalize (rounding errors)
 
-        whole_matrix_T = np.zeros_like(whole_matrix_R)
+    whole_matrix_R = whole_matrix/np.sum(whole_matrix,0)
 
-        whole_matrix = np.vstack([whole_matrix_R, whole_matrix_T])
+    whole_matrix_T = np.zeros_like(whole_matrix_R)
 
-        print(whole_matrix.shape)
+    whole_matrix = np.vstack([whole_matrix_R, whole_matrix_T])
 
-        A_matrix = np.zeros((1,int(len(angle_vector)/2)))
+    A_matrix = np.zeros((1,int(len(angle_vector)/2)))
 
-        allArray = COO(whole_matrix)
-        absArray = COO(A_matrix)
-        if save:
-            save_npz(savepath_RT, allArray)
-            save_npz(savepath_A, absArray)
+    allArray = COO(whole_matrix)
+    absArray = COO(A_matrix)
 
+    if save:
+        save_npz(savepath_RT, allArray)
+        save_npz(savepath_A, absArray)
 
     return allArray
 
@@ -87,53 +85,47 @@ def mirror_matrix(angle_vector, theta_intv, phi_intv, surf_name, options, struct
         savepath_RT = os.path.join(structpath, surf_name + front_or_rear + 'RT.npz')
         savepath_A = os.path.join(structpath, surf_name + front_or_rear + 'A.npz')
 
-    if os.path.isfile(savepath_RT) and save:
-        print('Existing angular redistribution matrices found')
-        allArray = load_npz(savepath_RT)
-
-    else:
-
-        if front_or_rear == "front":
+        if os.path.isfile(savepath_RT):
+            print('Existing angular redistribution matrices found')
+            allArray = load_npz(savepath_RT)
+            return allArray
 
-            angle_vector_th = angle_vector[:int(len(angle_vector)/2),1]
-            angle_vector_phi = angle_vector[:int(len(angle_vector)/2),2]
 
-            phis_out = fold_phi(angle_vector_phi + np.pi, options['phi_symmetry'])
+    if front_or_rear == "front":
 
+        angle_vector_th = angle_vector[:int(len(angle_vector)/2),1]
+        angle_vector_phi = angle_vector[:int(len(angle_vector)/2),2]
 
-        else:
-            angle_vector_th = angle_vector[int(len(angle_vector) / 2):, 1]
-            angle_vector_phi = angle_vector[int(len(angle_vector) / 2):, 2]
+        phis_out = fold_phi(angle_vector_phi + np.pi, options['phi_symmetry'])
 
-            phis_out = fold_phi(angle_vector_phi + np.pi, options['phi_symmetry'])
 
-        # matrix will be all zeros with just one '1' in each column/row. Just need to determine where it goes
+    else:
+        angle_vector_th = angle_vector[int(len(angle_vector) / 2):, 1]
+        angle_vector_phi = angle_vector[int(len(angle_vector) / 2):, 2]
 
-        binned_theta = np.digitize(angle_vector_th, theta_intv, right=True) - 1
+        phis_out = fold_phi(angle_vector_phi + np.pi, options['phi_symmetry'])
 
-        # print(binned_theta_out, theta_out, theta_intv)
+    # matrix will be all zeros with just one '1' in each column/row. Just need to determine where it goes
 
-        # print(binned_theta_in)
-        # print(binned_theta_out)
-        # -1 to give the correct index for the bins in phi_intv
+    binned_theta = np.digitize(angle_vector_th, theta_intv, right=True) - 1
 
+    bin_in = np.arange(len(angle_vector_phi))
 
-        bin_in = np.arange(len(angle_vector_phi))
+    phi_ind = [np.digitize(phi, phi_intv[binned_theta[i1]], right=True) - 1 for i1, phi in enumerate(phis_out)]
+    overall_bin = [np.argmin(abs(angle_vector[:,0] - binned_theta[i1])) + phi_i for i1, phi_i in enumerate(phi_ind)]
 
-        phi_ind = [np.digitize(phi, phi_intv[binned_theta[i1]], right=True) - 1 for i1, phi in enumerate(phis_out)]
-        overall_bin = [np.argmin(abs(angle_vector[:,0] - binned_theta[i1])) + phi_i for i1, phi_i in enumerate(phi_ind)]
+    whole_matrix = np.zeros((len(overall_bin)*2, len(overall_bin)))
 
-        whole_matrix = np.zeros((len(overall_bin)*2, len(overall_bin)))
+    whole_matrix[overall_bin, bin_in] = 1
 
-        whole_matrix[overall_bin, bin_in] = 1
 
+    A_matrix = np.zeros((1, len(overall_bin)))
 
-        A_matrix = np.zeros((1, len(overall_bin)))
+    allArray = COO(whole_matrix)
+    absArray = COO(A_matrix)
 
-        allArray = COO(whole_matrix)
-        absArray = COO(A_matrix)
-        if save:
-            save_npz(savepath_RT, allArray)
-            save_npz(savepath_A, absArray)
+    if save:
+        save_npz(savepath_RT, allArray)
+        save_npz(savepath_A, absArray)
 
     return allArray