Added checkpoint reader in rayleigh_diagnostics

post_processing/rayleigh_diagnostics.py

@@ -3266,3 +3266,18 @@ def streamfunction(vr,vt,r,cost,order=0):
 
     return psi
 
+###### Function for reading in the checkpoint files (i.e. P,PAB,T,TAB,W,WAB,Z,ZAB)
+def checkpoint_read(chk_file_string,nr,ntheta):
+    nell=(2*ntheta)//3
+    shape = (nell,nell)
+    i,j = np.indices(shape)
+    m = np.ceil(i <= j).astype(bool)
+    target_all = zeros((nell,nell,nr),dtype="complex")
+    length_half = int(len(fromfile(chk_file_string,"f8"))/2)
+    chunk_length = int(length_half/nr)
+    for i in range(nr):
+        target = np.zeros_like((m),dtype="complex")
+        target[m] = (fromfile(chk_file_string,"f8")[:length_half][i*chunk_length:(i+1)*chunk_length] 
+                + 1j*fromfile(chk_file_string,"f8")[length_half:][i*chunk_length:(i+1)*chunk_length])
+        target_all[:,:,i] = target
+    return target_all