update python support for movies with new file format

python/magic/TOreaders.py

@@ -6,11 +6,12 @@
 import numpy as np
 import matplotlib.pyplot as plt
 from .npfile import npfile
+from .movie import Movie
 from .log import MagicSetup
 from .libmagic import scanDir
 
 
-class TOMovie:
+class TOMovie(Movie):
     """
     This class allows to read and display the :ref:`TO_mov.TAG <secTO_movieFile>`
     generated when :ref:`l_TOmovie=.true. <varl_TOmovie>` is True.
@@ -22,11 +23,17 @@ class TOMovie:
     >>> t = TOMovie(file='TO_mov.N0m2', avg=True, levels=65, cm='seismic')
     """
 
-    def __init__(self, file=None, iplot=True, cm='seismic',
-                 cut=0.5, levels=33, avg=True, precision=np.float32):
+    def __init__(self, file=None, iplot=True, cm='seismic', datadir='.',
+                 cut=0.5, levels=33, avg=True, precision=np.float32,
+                 lastvar=None, nvar='all'):
         """
         :param file: the filename of the TO_mov file
         :type file: str
+        :param nvar: the number of frames of the movie file we want to plot
+                     starting from the last line
+        :type nvar: int
+        :param lastvar: the index of the last timesteps to be read
+        :type lastvar: int
         :param cm: the name of the color map
         :type cm: str
         :param levels: the number of contour levels
@@ -41,6 +48,8 @@ def __init__(self, file=None, iplot=True, cm='seismic',
         :param precision: precision of the input file, np.float32 for single
                           precision, np.float64 for double precision
         :type precision: str
+        :param datadir: working directory
+        :type datadir: str
         """
 
         if file is None:
@@ -58,53 +67,26 @@ def __init__(self, file=None, iplot=True, cm='seismic',
         else:
             filename = file
 
-        mot = re.compile(r'TO_mov\.(.*)')
-        end = mot.findall(filename)[0]
-
-        # DETERMINE THE NUMBER OF LINES BY READING THE LOG FILE
-        logfile = open('log.{}'.format(end), 'r')
-        mot = re.compile(r' ! WRITING TO MOVIE FRAME NO\s*(\d*).*')
-        for line in logfile.readlines():
-            if mot.match(line):
-                nlines = int(mot.findall(line)[0])
-        logfile.close()
+        filename = os.path.join(datadir, filename)
 
-        self.nvar = nlines
+        # Get the version
+        self._get_version(filename)
 
-        # READ the movie file
+        # Read the movie file
         infile = npfile(filename, endian='B')
-        # HEADER
-        version = infile.fort_read('|S64')
-        n_type, n_surface, const, n_fields = infile.fort_read(precision)
-        movtype = infile.fort_read(precision)
-        n_fields = int(n_fields)
-        if n_fields == 8:
+
+        # Header
+        self._read_header(infile, 0, precision)
+        if self.n_fields == 8:
             self.l_phase = True
         else:
             self.l_phase = False
-        self.movtype = np.asarray(movtype)
-        n_surface = int(n_surface)
-
-        # RUN PARAMETERS
-        runid = infile.fort_read('|S64')
-        n_r_mov_tot, n_r_max, n_theta_max, n_phi_tot, self.minc, self.ra, \
-            self.ek, self.pr, self.prmag, \
-            self.radratio, self.tScale = infile.fort_read(precision)
-        n_r_mov_tot = int(n_r_mov_tot)
-        self.n_r_max = int(n_r_max)
-        self.n_theta_max = int(n_theta_max)
-        self.n_phi_tot = int(n_phi_tot)
-
-        # GRID
-        self.radius = infile.fort_read(precision)
-        self.radius = self.radius[:self.n_r_max]  # remove inner core
-        rout = 1./(1.-self.radratio)
-        # rin = self.radratio/(1.-self.radratio)
-        self.radius *= rout
-        self.theta = infile.fort_read(precision)
-        self.phi = infile.fort_read(precision)
 
-        shape = (self.n_theta_max, self.n_r_max)
+        # Get the number of lines
+        self._get_nlines(datadir, filename, nvar, lastvar)
+
+        # Shape
+        self._get_data_shape(precision, allocate=False)
 
         self.time = np.zeros(self.nvar, precision)
         self.asVphi = np.zeros((self.nvar, self.n_theta_max, self.n_r_max),
@@ -118,28 +100,41 @@ def __init__(self, file=None, iplot=True, cm='seismic',
         if self.l_phase:
             self.penalty = np.zeros_like(self.asVphi)
 
-        # READ the data
+        # Read the data
+        # If one skip the beginning, nevertheless read but do not store
+        for i in range(self.var2-self.nvar):
+            if self.version == 2:
+                n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \
+                    movieDipLon, movieDipStrength, movieDipStrengthGeo = \
+                    infile.fort_read(precision)
+            else:
+                t_movieS = infile.fort_read(precision)
+            for ll in range(self.n_fields):
+                dat = infile.fort_read(precision, shape=self.shape, order='F')
         for k in range(self.nvar):
-            n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \
-                movieDipLon, movieDipStrength, movieDipStrengthGeo \
-                = infile.fort_read(precision)
+            if self.version == 2:
+                n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \
+                    movieDipLon, movieDipStrength, movieDipStrengthGeo \
+                    = infile.fort_read(precision)
+            else:
+                t_movieS = infile.fort_read(precision)
             self.time[k] = t_movieS
-            self.asVphi[k, ...] = infile.fort_read(precision, shape=shape,
+            self.asVphi[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                    order='F')
-            self.rey[k, ...] = infile.fort_read(precision, shape=shape,
+            self.rey[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                 order='F')
-            self.adv[k, ...] = infile.fort_read(precision, shape=shape,
+            self.adv[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                 order='F')
-            self.visc[k, ...] = infile.fort_read(precision, shape=shape,
+            self.visc[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                  order='F')
-            self.lorentz[k, ...] = infile.fort_read(precision, shape=shape,
+            self.lorentz[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                     order='F')
-            self.coriolis[k, ...] = infile.fort_read(precision, shape=shape,
+            self.coriolis[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                      order='F')
-            self.dtVp[k, ...] = infile.fort_read(precision, shape=shape,
+            self.dtVp[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                  order='F')
             if self.l_phase:
-                self.penalty[k, ...] = infile.fort_read(precision, shape=shape,
+                self.penalty[k, ...] = infile.fort_read(precision, shape=self.shape,
                                                         order='F')
 
         if iplot:

python/magic/__init__.py

@@ -23,7 +23,6 @@
         from .potExtra import *
         from .graph2vtk import *
         #from .graph2rst import *
-        from .movie3D import *
         from .spectralTransforms import *
         from .potential import *
     except ImportError as e:

python/magic/averages.py

@@ -35,7 +35,7 @@ def to_json(o, level=0):
             ret += SPACE * INDENT * (level + 1)
             ret += '"' + str(k) + '":' + SPACE
             ret += to_json(v, level + 1)
-  
+
         ret += NEWLINE + SPACE * INDENT * level + "}"
     elif isinstance(o, str):
         ret += '"' + o + '"'
@@ -74,12 +74,12 @@ class AvgField:
     This class computes the time-average properties from time series, spectra
     and radial profiles. It will store the input starting time in a small file
     named ``tInitAvg``, such that the next time you use it you don't need to
-    provide ``tstart`` again. By default, the outputs are stored in a 
+    provide ``tstart`` again. By default, the outputs are stored in a
     fully documented JSON file named avg.json: this is split into several
     categories, namely numerical parameters, physical parameters, time averaged
     scalar quantities, time averaged spectra and time-averaged radial profiles.
     The quantities stored in the JSON file are entirely controlled by an input
-    model file wich enlists the different quantities of interest. 
+    model file wich enlists the different quantities of interest.
     A default example file named ``model.json`` is provided in
     ``$MAGIC_HOME/python/magic``, and an example of how to build a dedicated one
     is provided below.
@@ -97,7 +97,7 @@ class AvgField:
                        'spectra': {},
                        'radial_profiles': {'powerR': ['viscDiss', 'buoPower']}
                      }
-    >>> # Compute the selected averages in the dirctory mydir                 
+    >>> # Compute the selected averages in the dirctory mydir
     >>> a = AvgField(datadir='mydir', model=json_model)
     """
 
@@ -409,7 +409,7 @@ class AvgStack:
     avg files to compute a global output which summarises all the outputs in one
     single file
 
-    >>> # Simply go through the directories listed in "runs.txt" and produce a 
+    >>> # Simply go through the directories listed in "runs.txt" and produce a
     >>> # local file named "my_results.json"
     >>> st = AvgStack(dirList="runs.txt", filename="my_results.json")
     >>> # Now also recompute each individual avg.json file in each directory
@@ -423,8 +423,8 @@ def __init__(self, dirList='runs.txt', filename='avg.json', datadir='.',
         :param dirList: the filename of the list of directories, lines starting with
                         a hash are omitted by the reader
         :type dirList: str
-        :param filename: 
-        :param recompute: recompute each individual average file in each single 
+        :param filename:
+        :param recompute: recompute each individual average file in each single
                           directory when set to True
         :type recompute: bool
         :param datadir: working directory

python/magic/checker.py

@@ -22,7 +22,7 @@ def MagicCheck(tstart=None):
     if the power.TAG and some spectra have been produced in the current directory.
     If in addition the tInitAvg file is also there in the directory it averages
     only from this starting time.
-    
+
     >>> MagicCheck(tstart=10.)
     """
 
@@ -48,7 +48,7 @@ def MagicCheck(tstart=None):
     ohmDiss_avg = avgField(ts.time[mask], ts.ohmDiss[mask])
     viscDiss_avg = avgField(ts.time[mask], ts.viscDiss[mask])
     ratio = 100*abs(buoPower_avg+buoPower_chem_avg+ohmDiss_avg+viscDiss_avg) / \
-            (buoPower_avg+buoPower_chem_avg) 
+            (buoPower_avg+buoPower_chem_avg)
 
     print(bcolors.BOLD + bcolors.UNDERLINE + 'Power balance:' + bcolors.ENDC)
     print('Power injected   : {:.5e}'.format(buoPower_avg+buoPower_chem_avg))

python/magic/checkpoint.py

@@ -58,7 +58,7 @@ def get_map(lm_max, lmax, mmin, mmax, minc):
 def interp_one_field(field, rold, rnew, rfac=None):
     """
     This routine interpolates a complex input field from an old radial grid
-    to a new one. 
+    to a new one.
 
     :param field: the field to be interpolated
     :type field: numpy.ndarray
@@ -164,7 +164,7 @@ def __init__(self, l_read=True, filename=None, endian='l'):
 
     def read(self, filename, endian='l'):
         """
-        This routine is used to read a checkpoint file. 
+        This routine is used to read a checkpoint file.
 
         :param filename: name of the checkpoint file
         :type filename: str
@@ -179,7 +179,7 @@ def read(self, filename, endian='l'):
         self.version = np.fromfile(file, fmt, count=1)[0]
         fmt = '{}f8'.format(prefix)
         self.time = np.fromfile(file, fmt, count=1)[0]
-        
+
         # Time scheme
         self.tscheme_family = file.read(10).decode()
         nexp, nimp, nold = np.fromfile(file, dtype=np.int32, count=3)
@@ -574,7 +574,7 @@ def xshells2magic(self, xsh_trailing, n_r_max, rscheme='cheb',
 
         # When nalias=60 ntheta=lmax: trick to have lmax in MagIC's header
         self.nalias = 60
-        self.n_theta_max = self.l_max 
+        self.n_theta_max = self.l_max
         self.n_phi_tot = self.l_max
         self.n_r_ic_max = 1
 
@@ -727,7 +727,7 @@ def graph2rst(self, gr, filename='checkpoint_ave.from_chk'):
         s3D = rr3D*np.sin(th3D)
         omr = 1./s3D*(thetaderavg(np.sin(th3D)*gr.vphi, order=4) -
                       phideravg(gr.vtheta, minc=self.minc))
- 
+
         for i in range(self.n_r_max):
             om = sh.spat_spec(omr[:, :, i])
             self.ztor[i, 1:] = om[1:]/(sh.ell[1:]*(sh.ell[1:]+1)) * \
@@ -777,7 +777,7 @@ def graph2rst(self, gr, filename='checkpoint_ave.from_chk'):
                     rr = self.radius_ic[i-1]
                     rdep[1:] = (self.radius_ic[i-1]/ri)**(sh.ell[1:]+1)
                 self.bpol_ic[i, 1:] = vr[1:]/(sh.ell[1:]*(sh.ell[1:]+1))*rr**2
-                # Not stable: 
+                # Not stable:
                 #if self.radius_ic[i] >= 0.01:
                 #    mask = ( self.lm2l <= 2 ) * ( self.lm2m <= 2)
                 #    self.bpol_ic[i, mask] /= rdep[mask]

python/magic/cyl.py

@@ -96,7 +96,7 @@ def sph2cyl_plane(data, rad, ns):
     def zavg(input, radius, ns, minc, save=True, filename='vp.pickle',
              normed=True, colat=None):
         """
-        This function computes a z-integration of a list of input arrays 
+        This function computes a z-integration of a list of input arrays
         (on the spherical grid). This works well for 2-D (phi-slice) arrays.
         In case of 3-D arrays, only one element is allowed (too demanding
         otherwise).
@@ -114,7 +114,7 @@ def zavg(input, radius, ns, minc, save=True, filename='vp.pickle',
         :type save: bool
         :param filename: name of the output pickle when save=True
         :type filename: str
-        :param normed: a boolean to specify if ones wants to simply integrate 
+        :param normed: a boolean to specify if ones wants to simply integrate
                        over z or compute a z-average (default is True: average)
         :type normed: bool
         :param colat: an optional array containing the colatitudes
@@ -186,8 +186,8 @@ def zavg(input, radius, ns, minc, save=True, filename='vp.pickle',
     def zavg(input, radius, ns, minc, save=True, filename='vp.pickle',
              normed=True):
         """
-        This function computes a z-integration of a list of input arrays 
-        (on the spherical grid). This works well for 2-D (phi-slice) 
+        This function computes a z-integration of a list of input arrays
+        (on the spherical grid). This works well for 2-D (phi-slice)
         arrays. In case of 3-D arrays, only one element is allowed
         (too demanding otherwise).
 

python/magic/graph.py

@@ -46,7 +46,7 @@ def getGraphEndianness(filename):
     except ValueError:
         f = open(filename, 'rb')
         # Little endian
-        version = np.fromfile(f, 'i4', count=1)[0]
+        version = np.fromfile(f, np.int32, count=1)[0]
         endian = 'l'
         if abs(version) > 100:
             f.close()

python/magic/libmelt.py

@@ -183,7 +183,7 @@ def stack_two(self, theta_new, time_new, rmelt_new):
         # If the number of latitudinal grid points has changed, one
         # needs to interpolate
         if len(self.colatitude) != len(theta_new):
-            it = interp1d(self.colatitude, self.rmelt, axis=-1, 
+            it = interp1d(self.colatitude, self.rmelt, axis=-1,
                           fill_value='extrapolate')
             self.rmelt = it(theta_new)
             self.colatitude = theta_new
@@ -208,7 +208,7 @@ def plot(self, levels=17, cm='viridis'):
         fig = plt.figure()
         ax = fig.add_subplot(111)
 
-        im = ax.contourf(self.time, self.colatitude, self.rmelt.T, levels, 
+        im = ax.contourf(self.time, self.colatitude, self.rmelt.T, levels,
                          cmap=plt.get_cmap(cm))
         ax.set_xlabel('Time')
         ax.set_ylabel('Colatitude')