Allow separate method for loading

Allow loading like this:

snap = pa.Snapshot('data/small_non_comoving_007.hdf5')

which is a bit simpler than the more verbose

s = pa.Snapshot(basedir='data', basename='small_non_comoving', snapnum=7)

especially when working in the terminal

paicos/readers/arepo_snap.py

@@ -56,7 +56,7 @@ class Snapshot(PaicosReader):
 
     # pylint: disable=too-many-instance-attributes
 
-    def __init__(self, basedir, snapnum, basename="snap", load_all=False,
+    def __init__(self, basedir, snapnum=None, basename="snap", load_all=False,
                  to_physical=False, load_catalog=None, verbose=False):
         """
         Initialize the Snapshot class.
@@ -86,6 +86,8 @@ def __init__(self, basedir, snapnum, basename="snap", load_all=False,
                          load_all=load_all, to_physical=to_physical,
                          basesubdir='snapdir', verbose=verbose)
 
+        assert self.snapnum is not None
+
         self.load_catalog = load_catalog
 
         if not hasattr(self, "dic_selection_index"):

paicos/readers/paicos_readers.py

@@ -49,56 +49,66 @@ def __init__(self, basedir='.', snapnum=None, basename="snap",
 
         verbose (bool): whether to print information, default is False
         """
-
-        self.basedir = basedir
-        self.snapnum = snapnum
-        self.basename = basename
-
         assert to_physical is False, 'to_physical not yet implemented!'
 
         self.to_physical = to_physical
-        self.basesubdir = basesubdir
         self.verbose = verbose
 
-        # Add a forward slash at the end of basedir if not already present
-        if basedir[-1] != '/':
-            basedir += '/'
-
-        # If snapnum is None, filename is basedir + basename + '.hdf5'
-        if snapnum is None:
-            self.filename = basedir + basename + '.hdf5'
-            err_msg = f'File: {self.filename} not found'
-            if not os.path.exists(self.filename):
-                raise FileNotFoundError(err_msg)
+        if '.hdf5' in basedir:
+            # User is trying to load a single hdf5 file directly, let's see if
+            # that can work
+            self.filename = str(basedir)
+            self.multi_file = False
+            basedir, basename, snapnum = util._split_filename(self.filename)
+            self.basedir = basedir
+            self.snapnum = snapnum
+            self.basename = basename
         else:
-            # Set filenames for single_file and multi_file cases
-            single_file = basename + "_{:03d}.hdf5"
-            multi_file = basesubdir + '_{:03d}/' + basename + '_{:03d}.{}.hdf5'
-            multi_wo_dir = basename + '_{:03d}.{}.hdf5'
-
-            single_file = basedir + single_file.format(snapnum)
-            multi_file = basedir + multi_file.format(snapnum, snapnum, '{}')
-            multi_wo_dir = basedir + multi_wo_dir.format(snapnum, '{}')
-
-            # Check if single_file exists
-            if os.path.exists(single_file):
-                self.multi_file = False
-                self.filename = single_file
-            elif os.path.exists(multi_file.format(0)):
-                self.multi_file = True
-                self.first_file_name = self.filename = multi_file.format(0)
-                self.multi_filename = multi_file
-                self.no_subdir = False
-            elif os.path.exists(multi_wo_dir.format(0)):
-                self.multi_file = True
-                self.first_file_name = self.filename = multi_wo_dir.format(0)
-                self.multi_filename = multi_file
-                self.no_subdir = True
+            self.basedir = basedir
+            self.snapnum = snapnum
+            self.basename = basename
+
+            self.basesubdir = basesubdir
+
+            # Add a forward slash at the end of basedir if not already present
+            if basedir[-1] != '/':
+                basedir += '/'
+
+            # If snapnum is None, filename is basedir + basename + '.hdf5'
+            if snapnum is None:
+                self.filename = basedir + basename + '.hdf5'
+                err_msg = f'File: {self.filename} not found'
+                if not os.path.exists(self.filename):
+                    raise FileNotFoundError(err_msg)
             else:
-                err_msg = "File not found. Tried locations:\n{}\n{}\n{}"
-                err_msg = err_msg.format(single_file, multi_file.format(0),
-                                         multi_wo_dir.format(0))
-                raise FileNotFoundError(err_msg)
+                # Set filenames for single_file and multi_file cases
+                single_file = basename + "_{:03d}.hdf5"
+                multi_file = basesubdir + '_{:03d}/' + basename + '_{:03d}.{}.hdf5'
+                multi_wo_dir = basename + '_{:03d}.{}.hdf5'
+
+                single_file = basedir + single_file.format(snapnum)
+                multi_file = basedir + multi_file.format(snapnum, snapnum, '{}')
+                multi_wo_dir = basedir + multi_wo_dir.format(snapnum, '{}')
+
+                # Check if single_file exists
+                if os.path.exists(single_file):
+                    self.multi_file = False
+                    self.filename = single_file
+                elif os.path.exists(multi_file.format(0)):
+                    self.multi_file = True
+                    self.first_file_name = self.filename = multi_file.format(0)
+                    self.multi_filename = multi_file
+                    self.no_subdir = False
+                elif os.path.exists(multi_wo_dir.format(0)):
+                    self.multi_file = True
+                    self.first_file_name = self.filename = multi_wo_dir.format(0)
+                    self.multi_filename = multi_file
+                    self.no_subdir = True
+                else:
+                    err_msg = "File not found. Tried locations:\n{}\n{}\n{}"
+                    err_msg = err_msg.format(single_file, multi_file.format(0),
+                                             multi_wo_dir.format(0))
+                    raise FileNotFoundError(err_msg)
 
         with h5py.File(self.filename, 'r') as f:
             self.Header = dict(f['Header'].attrs)

paicos/util.py

@@ -42,6 +42,43 @@ def get_project_root_dir():
 root_dir = get_project_root_dir()
 
 
+def _split_filename(filenamewithpath):
+    """
+    A convenience function for getting the basedir, basename and snapsnum
+    from a filename including its path.
+
+    Example usage:
+        print(split_filename('data/small_non_comoving_007.hdf5'))
+        print(split_filename('small_non_comoving.hdf5'))
+        print(split_filename('small_non_comoving.0.hdf5'))
+
+    """
+    err_msg = "needs to be a hdf5 file, .e.g. filenamewithpath = 'data/snap_010.hdf5'"
+
+    assert filenamewithpath[-5:] == '.hdf5', err_msg
+
+    snapnum = None
+
+    index = filenamewithpath[::-1].find('_')
+    if index != -1:
+        underscore_loc = -index - 1
+
+    if filenamewithpath[underscore_loc] == '_':
+        snapnum_str = filenamewithpath[underscore_loc + 1:underscore_loc + 4]
+        if np.char.isdigit(snapnum_str):
+            snapnum = int(snapnum_str)
+
+    filename = os.path.basename(filenamewithpath)
+    basedir = filenamewithpath.replace(filename, '')
+
+    if snapnum is None:
+        basename = filename[:-5]
+    else:
+        basename = filename[:underscore_loc]
+
+    return basedir, basename, snapnum
+
+
 def save_dataset(hdf5file, name, data=None, data_attrs={},
                  group=None, group_attrs={}):
     """