anta-database 0.1.4__tar.gz → 0.2.0__tar.gz

{anta_database-0.1.4 → anta_database-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anta_database
-Version: 0.1.4
+Version: 0.2.0
 Summary: SQLite database for the AntArchitecture Community Data
 Author: Antoine Hermant
 Author-email: antoine.hermant@etik.com
@@ -14,12 +14,14 @@ Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: Topic :: Scientific/Engineering :: Physics
 Classifier: Programming Language :: Python
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pandas
 Requires-Dist: pyproj
 Requires-Dist: matplotlib
 Requires-Dist: colormaps
+Requires-Dist: tqdm
 Dynamic: license-file
 
 <h1 align="center">
@@ -27,6 +29,7 @@ Dynamic: license-file
 </h1><br>
 
 [![PyPi](https://img.shields.io/pypi/v/anta_database)](https://pypi.org/project/anta_database/)
+[![Conda](https://img.shields.io/conda/vn/conda-forge/anta_database)](https://anaconda.org/conda-forge/anta_database)
 [![Downloads](https://img.shields.io/pypi/dm/anta_database)](https://pypi.org/project/anta_database)
 
 [![GitHub issues](https://img.shields.io/badge/issue_tracking-github-blue.svg)](https://github.com/antoinehermant/anta_database/issues)

{anta_database-0.1.4 → anta_database-0.2.0}/README.md

@@ -3,6 +3,7 @@
 </h1><br>
 
 [![PyPi](https://img.shields.io/pypi/v/anta_database)](https://pypi.org/project/anta_database/)
+[![Conda](https://img.shields.io/conda/vn/conda-forge/anta_database)](https://anaconda.org/conda-forge/anta_database)
 [![Downloads](https://img.shields.io/pypi/dm/anta_database)](https://pypi.org/project/anta_database)
 
 [![GitHub issues](https://img.shields.io/badge/issue_tracking-github-blue.svg)](https://github.com/antoinehermant/anta_database/issues)

{anta_database-0.1.4 → anta_database-0.2.0}/anta_database/__init__.py

@@ -1,3 +1,4 @@
 from .database import *
 from .indexing import *
 from .compiling import *
+from .plotting import *

{anta_database-0.1.4 → anta_database-0.2.0}/anta_database/anta_database.py

@@ -1,3 +1,4 @@
 from database import *
 from indexing import *
 from compiling import *
+from plotting import *

anta_database-0.2.0/anta_database/compiling/compile.py

@@ -0,0 +1,262 @@
+import os
+import time
+import pandas as pd
+import numpy as np
+import glob
+from pyproj import Transformer
+from typing import Union
+from multiprocessing import Pool, cpu_count
+from tqdm import tqdm
+
+class CompileDatabase:
+    def __init__(self, dir_list: Union[str, list[str]], file_type: str = 'layer', wave_speed: Union[None, float] = None, firn_correction: Union[None, float] = None, comp: bool = True, post: bool = True) -> None:
+        self.dir_list = dir_list
+        self.wave_speed = wave_speed
+        self.firn_correction = firn_correction
+        self.file_type = file_type
+        self.comp = comp
+        self.post = post
+
+    def get_dict_ages(self, tab_file) -> dict:
+        ages = pd.read_csv(tab_file, header=None, sep='\t', names=['file', 'age'])
+        return dict(zip(ages['file'], ages['age']))
+
+    def _pre_compile_checks(self, dir_list: list[str]) -> bool:
+        missing = False
+        for dir_path in dir_list:
+            raw_dir = f"{dir_path}/raw/"
+            if not os.path.exists(raw_dir):
+                print(f"{raw_dir} does not exist")
+                missing = True
+        return not missing
+
+    def compile(self, cpus: int = cpu_count()-1) -> None:
+        if not isinstance(self.dir_list, list):
+            self.dir_list = [self.dir_list]
+
+        start_time = time.time()
+        if not self._pre_compile_checks(self.dir_list):
+            return
+
+        if self.comp is True:
+            all_files_list = []
+            for dir_ in self.dir_list:
+                files = glob.glob(f"{dir_}/raw/*.*")
+                for file_path in files:
+                    all_files_list.append({
+                        'dir_path': dir_,
+                        'file': os.path.basename(file_path),
+                        'file_path': file_path
+                    })
+
+            num_tasks = len(all_files_list)
+            num_workers = min(num_tasks, cpus)
+
+            print('\n',
+                    'Will start compiling', num_tasks, 'raw files\n'
+                    '\n   ', num_workers, 'worker(s) allocated out of', cpu_count(), 'available cpus\n')
+
+            if num_workers > 1:
+                with Pool(num_workers) as pool:
+                    for _ in tqdm(pool.imap_unordered(self._compile, all_files_list), total=num_tasks):
+                        pass
+            else:
+                for file_dict in tqdm(all_files_list, desc="Processing"):
+                    self._compile(file_dict=file_dict)
+
+        if self.post is True:
+            all_dirs = []
+            for dir_ in self.dir_list:
+                dirs = [d for d in glob.glob(f"{dir_}/pkl/*") if os.path.isdir(d)]
+                all_dirs.extend(dirs)
+
+            num_tasks = len(all_dirs)
+            num_workers = min(num_tasks, cpus)
+
+            print('\n',
+                    'Will start post compilation of', len(all_dirs), 'traces\n'
+                    '\n   ', num_workers, 'worker(s) allocated out of', cpu_count(), 'available cpus\n')
+
+            if num_workers > 1:
+                with Pool(num_workers) as pool:
+                    for _ in tqdm(pool.imap_unordered(self._post_compilation, all_dirs), total=num_tasks):
+                        pass
+            else:
+                for trace_dir in tqdm(all_dirs, desc='Processing'):
+                    self._post_compilation(trace_dir=trace_dir)
+
+        elapsed = time.time() - start_time
+        print(f"\nCompilation completed in {elapsed:.2f} seconds")
+
+    def _compile(self, file_dict) -> None:
+
+        _, ext = os.path.splitext(file_dict['file'])
+        ages = self.get_dict_ages(f'{file_dict['dir_path']}/IRH_ages.tab')
+        original_new_columns = pd.read_csv(f'{file_dict['dir_path']}/original_new_column_names.csv')
+
+        if ext == '.tab':
+            sep='\t'
+        elif ext == '.csv':
+            sep=','
+        else:
+            print(f"{ext}: File type not supported...")
+            return
+
+        ds = pd.read_csv(file_dict['file_path'], comment="#", header=0, sep=sep, na_values=['-9999', 'NaN', ''])
+        _, file_name = os.path.split(file_dict['file_path'])
+        file_name_, ext = os.path.splitext(file_name)
+
+        ds = ds[ds.columns.intersection(original_new_columns.columns)]
+        ds.columns = original_new_columns[ds.columns].iloc[0].values  # renaming the columns
+
+        if 'IceThk' in ds.columns and 'SurfElev' in ds.columns and not 'BedElev' in ds.columns:
+            ds['BedElev'] = ds['SurfElev'] - ds['IceThk']
+        if 'IceThk' in ds.columns and 'BedElev' in ds.columns and not 'SurfElev' in ds.columns:
+            ds['SurfElev'] = ds['BedElev'] + ds['IceThk']
+        if 'SurfElev' in ds.columns and 'BedElev' in ds.columns and not 'IceThk' in ds.columns:
+            ds['IceThk'] = ds['SurfElev'] - ds['BedElev']
+
+        if self.wave_speed:
+            for var in ['IceThk', 'BedElev']:
+                if var in ds.columns:
+                    ds[var] *= self.wave_speed
+        if self.firn_correction:
+            for var in ['IceThk', 'BedElev']:
+                if var in ds.columns:
+                    ds[var] += self.firn_correction
+
+        if 'x' not in ds.columns and 'y' not in ds.columns:
+            if 'lon' in ds.columns and 'lat' in ds.columns:
+                transformer = Transformer.from_proj(
+                    "EPSG:4326",  # source: WGS84 (lon/lat)
+                    "+proj=stere +lon_0=0 +lat_0=-90 +lat_ts=-71 +datum=WGS84 +units=m +no_defs",  # target: polar
+                    always_xy=True
+                )
+                ds['x'], ds['y'] = transformer.transform(ds['lon'].values, ds['lat'].values)
+        elif 'lon' not in ds.columns and 'lat' not in ds.columns:
+            if 'x' in ds.columns and 'y' in ds.columns:
+                inverse_transformer = Transformer.from_proj(
+                    "+proj=stere +lon_0=0 +lat_0=-90 +lat_ts=-71 +datum=WGS84 +units=m +no_defs",  # source: polar
+                    "EPSG:4326",  # target: WGS84 (lon/lat)
+                    always_xy=True
+                )
+                ds['lon'], ds['lat'] = inverse_transformer.transform(ds['x'].values, ds['y'].values)
+        elif 'lon' in ds.columns and 'lat' in ds.columns and 'x' in ds.columns and 'y' in ds.columns:
+            pass
+        else:
+            print('No coordinates found in the dataset')
+            return
+
+        if self.file_type == 'layer':
+            age = str(ages[file_name_])
+            if self.wave_speed:
+                ds['IRHDepth'] *= self.wave_speed
+            if self.firn_correction:
+                ds['IRHDepth'] += self.firn_correction
+
+            ds['Trace_ID'] = ds['Trace_ID'].astype(str)
+            ds['Trace_ID'] = ds['Trace_ID'].str.replace(r'/\s+', '_') # Replace slashes with underscores, otherwise the paths can get messy
+            ds['Trace_ID'] = ds['Trace_ID'].str.replace('/', '_')
+
+            ds.set_index('Trace_ID', inplace=True)
+
+            for trace_id in np.unique(ds.index):
+                ds_trace = ds.loc[trace_id].copy()
+                ds_trace = ds_trace.drop_duplicates(subset=['x', 'y']) # Some datasets showed duplicated data
+                if 'distance' not in ds_trace.columns:
+                    x = ds_trace[['x', 'y']]
+                    distances = np.sqrt(np.sum(np.diff(x, axis=0)**2, axis=1))
+                    cumulative_distance = np.concatenate([[0], np.cumsum(distances)])
+                    ds_trace['distance'] = cumulative_distance
+                elif 'Distance [km]' in original_new_columns.columns:
+                    ds_trace['distance'] *= 1000 # if distance in km, convert to meters
+
+                if age in ['IceThk', 'BedElev', 'SurfElev', 'BasalUnit']:
+                    ds_trace = ds_trace.rename(columns={'IRHDepth': age})
+                    ds_trace_file = f'{file_dict['dir_path']}/pkl/{trace_id}/{age}.pkl' # if var instead of age, call the file as var.pkl
+                else:
+                    ds_trace_file = f'{file_dict['dir_path']}/pkl/{trace_id}/{file_name_}.pkl' # else use the same file name.pkl
+
+                os.makedirs(f'{file_dict['dir_path']}/pkl/{trace_id}' , exist_ok=True)
+                ds_trace.to_pickle(ds_trace_file)
+
+        elif self.file_type == 'trace':
+            if 'distance' not in ds.columns:
+                x = ds[['x', 'y']]
+                distances = np.sqrt(np.sum(np.diff(x, axis=0)**2, axis=1))
+                cumulative_distance = np.concatenate([[0], np.cumsum(distances)])
+                ds['distance'] = cumulative_distance
+            elif 'Distance [km]' in original_new_columns.columns:
+                ds['distance'] *= 1000 # if distance in km, convert to meters
+
+            trace_id = file_name_
+            os.makedirs(f'{file_dict['dir_path']}/pkl/{trace_id}' , exist_ok=True)
+
+            ages = {key: ages[key] for key in ds.columns if key in ages}
+
+            for IRH in ages:
+                age = str(ages.get(IRH))
+                ds_IRH = ds[IRH]
+                ds_IRH = pd.DataFrame({
+                    'lon': ds['lon'],
+                    'lat': ds['lat'],
+                    'x': ds['x'],
+                    'y': ds['y'],
+                    'distance': ds['distance'],
+                    'IRHDepth': ds_IRH,
+                })
+                if self.wave_speed:
+                    ds_IRH['IRHDepth'] *= self.wave_speed
+                if self.firn_correction:
+                    ds_IRH['IRHDepth'] += self.firn_correction
+
+                for var in ['IceThk', 'BedElev', 'SurfElev', 'BasalUnit']:
+                    if var in ds.columns:
+                        ds_IRH[var] = ds[var]
+
+                ds_trace_file = f'{file_dict['dir_path']}/pkl/{trace_id}/{IRH}.pkl'
+                ds_IRH.to_pickle(ds_trace_file)
+
+    def compute_irh_density(self, trace_dir) -> None:
+        files = glob.glob(f"{trace_dir}/*.pkl")
+
+        unwanted = {'IceThk.pkl', 'SurfElev.pkl', 'BasalUnit.pkl', 'BedElev.pkl', 'IRHDensity.pkl'}
+        files = [f for f in glob.glob(f"{trace_dir}/*.pkl") if os.path.basename(f) not in unwanted]
+        if len(files) > 1:
+            dfs = [pd.read_pickle(f) for f in files]
+            dfs = pd.concat(dfs)
+        elif len(files) == 1:
+            dfs = pd.read_pickle(files[0])
+        else:
+            return
+
+        dfs = dfs[['x','y','IRHDepth']]
+        valid = dfs.dropna(subset=['IRHDepth'])
+        density = valid.groupby(['x', 'y']).size().reset_index(name='IRHDensity')
+
+        density_file = f'{trace_dir}/IRHDensity.pkl'
+        density.to_pickle(density_file)
+
+    def extract_vars(self, trace_dir: str) -> None:
+        files = glob.glob(f"{trace_dir}/*.pkl")
+
+        unwanted = {'IceThk.pkl', 'SurfElev.pkl', 'BasalUnit.pkl', 'BedElev.pkl', 'IRHDensity.pkl'}
+        files = [f for f in glob.glob(f"{trace_dir}/*.pkl") if os.path.basename(f) not in unwanted]
+        if len(files) > 1:
+            dfs = [pd.read_pickle(f) for f in files]
+            dfs = pd.concat(dfs).drop_duplicates(subset=['x', 'y'])
+        elif len(files) == 1:
+            dfs = pd.read_pickle(files[0])
+        else:
+            return
+
+        for var in ['IceThk', 'BedElev', 'SurfElev']:
+            if var in dfs.columns:
+                ds_var = dfs[['x', 'y', 'distance', var]]
+                var_file = f'{trace_dir}/{var}.pkl'
+                ds_var.to_pickle(var_file)
+
+    def _post_compilation(self, trace_dir: str) -> None:
+        self.extract_vars(trace_dir)
+        self.compute_irh_density(trace_dir)
+

{anta_database-0.1.4 → anta_database-0.2.0}/anta_database/database/database.py

@@ -1,21 +1,18 @@
-import os, sys
+import os
 import sqlite3
 import pandas as pd
 import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.colors import LinearSegmentedColormap
-import colormaps as cmaps
-from contextlib import contextmanager
 from typing import Union, List, Dict, Tuple
 
+from anta_database.plotting.plotting import Plotting
+
 class Database:
-    def __init__(self, database_dir: str, file_db: str = 'AntADatabase.db'):
-        module_dir = os.path.dirname(os.path.abspath(__file__))
-        self.gl = pd.read_pickle(os.path.join(module_dir, 'plotting', 'GL.pkl'))
+    def __init__(self, database_dir: str, file_db: str = 'AntADatabase.db') -> None:
         self.db_dir = database_dir
         self.file_db = file_db
         self.file_db_path = os.path.join(self.db_dir, file_db)
         self.md = None
+        self._plotting = None
 
     def _build_query_and_params(self, age: Union[None, str, List[str]]=None, var: Union[None, str, List[str]]=None, author: Union[None, str, List[str]]=None, line: Union[None, str, List[str]]=None, select_clause='') -> Tuple[str, List[Union[str, int]]]:
         """
@@ -29,7 +26,6 @@ class Database:
         '''
         conditions = []
         params = []
-
         for field, column in [
             (age, 'd.age'),
             (var, 'd.var'),
@@ -38,17 +34,31 @@ class Database:
         ]:
             if field is not None:
                 if isinstance(field, list):
-                    placeholders = ','.join(['?'] * len(field))
-                    conditions.append(f'{column} IN ({placeholders})')
-                    params.extend(field)
+                    # For lists, use IN for exact matches, or LIKE for wildcards
+                    like_conditions = []
+                    in_values = []
+                    for item in field:
+                        if '%' in item:
+                            like_conditions.append(f"{column} LIKE ?")
+                            params.append(item)
+                        else:
+                            in_values.append(item)
+                    if like_conditions:
+                        conditions.append('(' + ' OR '.join(like_conditions) + ')')
+                    if in_values:
+                        placeholders = ','.join(['?'] * len(in_values))
+                        conditions.append(f"{column} IN ({placeholders})")
+                        params.extend(in_values)
                 else:
-                    conditions.append(f'{column} = ?')
+                    # For single values, use = or LIKE
+                    if '%' in field:
+                        conditions.append(f"{column} LIKE ?")
+                    else:
+                        conditions.append(f"{column} = ?")
                     params.append(field)
-
         if conditions:
             query += ' WHERE ' + ' AND '.join(conditions)
-
-        query += 'ORDER BY CAST(d.age AS INTEGER) ASC'
+        query += ' ORDER BY CAST(d.age AS INTEGER) ASC'
         return query, params
 
     def _get_file_metadata(self, file_path) -> Dict:
@@ -207,179 +217,12 @@ class Database:
             metadata = self._get_file_metadata(file_path)
             yield df, metadata
 
-    def custom_cmap(self):
-
-        cm1 = cmaps.torch_r
-        cm2 = cmaps.deep_r
-
-        cm1_colors = cm1(np.linspace(0.05, 0.8, 256))  # Shape: (256, 4) RGBA
-        cm2_colors = cm2(np.linspace(0.1, 1, 256))  # Shape: (256, 4) RGBA
-        combined_colors = np.vstack((cm1_colors, cm2_colors))
-        custom_cmap = LinearSegmentedColormap.from_list(
-            'custom_cmap',  # Name your colormap
-            combined_colors,
-            N=512  # Double the resolution (256 + 256)
-        )
-        return custom_cmap
-
-    def is_notebook(self) -> bool:
-        from IPython import get_ipython
-        try:
-            shell = get_ipython().__class__.__name__
-            if shell == 'ZMQInteractiveShell':
-                return True   # Jupyter notebook or qtconsole
-            elif shell == 'TerminalInteractiveShell':
-                return False  # Terminal running IPython
-            else:
-                return False  # Other type (?)
-        except NameError:
-            return False      # Probably standard Python interpreter
-
-    @contextmanager
-    def plot_context(self, close=None):
-        if close is None:
-            close = not self.is_notebook()
-        try:
-            yield
-        finally:
-            if close:
-                plt.close()
-
-    def plotXY(self,
-               metadata: Union[None, Dict, 'MetadataResult'] = None,
-               downscale_factor: Union[None, int] = None,
-               title: str = '',
-               xlim: tuple = (None, None),
-               ylim: tuple = (None, None),
-               scale_factor: float = 1.0,
-               latex:bool = False,
-               save: Union[str, None] = None) -> None:
-
-
-        if metadata:
-            metadata = metadata
-        elif self.md:
-            metadata = self.md
-        else:
-            print('Please provide metadata of the files you want to generate the data from...')
-            return
-
-        cmaps = self.custom_cmap()
-        authors = list(metadata['author'])
-        color_indices = np.linspace(0.1, 0.9, len(authors))
-        colors = {author: cmaps(i) for i, author in zip(color_indices, authors)}
-
-        if latex:
-            from matplotlib import rc
-            rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
-            rc('text', usetex=True)
-
-        plt.figure()
-        plt.plot(self.gl.x/1000, self.gl.y/1000, linewidth=1, color='k')
-        for df, md in self.data_generator(metadata, downscale_factor=downscale_factor):
-            author = md['author']
-            plt.plot(df.x/1000, df.y/1000, linewidth=0.8, color=colors[author])
-
-        ax = plt.gca()
-        x0, x1 = ax.get_xlim() if xlim == (None, None) else xlim
-        y0, y1 = ax.get_ylim() if ylim == (None, None) else ylim
-
-        for author in authors:
-            citation = self.query(author=author)['reference']
-            plt.plot([], [], color=colors[author], label=citation)
-        plt.legend(loc='lower left', fontsize=16*scale_factor)
-
-        x_extent = x1 - x0
-        y_extent = y1 - y0
+    @property
+    def plot(self):
+        if self._plotting is None:
+            self._plotting = Plotting(self)
+        return self._plotting
 
-        aspect_ratio = y_extent / x_extent
-
-        # Set the figure size based on the aspect ratio and a scale factor
-        plt.gcf().set_size_inches(8 * scale_factor, 8 * aspect_ratio * scale_factor)
-        ax.set_xlabel('x [km]')
-        ax.set_ylabel('y [km]')
-        ax.set_xlim(*xlim)
-        ax.set_ylim(*ylim)
-        # self.scale_axes(ax, factor=1000)
-        plt.title(title, fontsize=30*scale_factor)
-        plt.tight_layout()
-
-        with self.plot_context():
-
-            if save:
-                plt.savefig(save, dpi=200)
-                print('Figure saved as', save)
-            else:
-                plt.show()
-
-    def plotXYDepth(self,
-                    metadata: Union[None, Dict, 'MetadataResult'] = None,
-                    downscale_factor: Union[None, int] = None,
-                    title: str = '',
-                    xlim: tuple = (None, None),
-                    ylim: tuple = (None, None),
-                    scale_factor: float = 1.0,
-                    latex:bool = False,
-                    cmaps = cmaps.torch_r,
-                    save: Union[str, None] = None) -> None:
-
-        if metadata:
-            metadata = metadata
-        elif self.md:
-            metadata = self.md
-        else:
-            print('Please provide metadata of the files you want to generate the data from...')
-            return
-
-        if latex:
-            from matplotlib import rc
-            rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})
-            rc('text', usetex=True)
-
-        fig, ax = plt.subplots(figsize=(10,6))
-        plt.plot(self.gl.x/1000, self.gl.y/1000, linewidth=1, color='k')
-        for df, md in self.data_generator(metadata, downscale_factor=downscale_factor):
-            if md['age']:
-                var = md['age']
-                label = f'{var} isochrone depth [m]'
-            elif md['var']:
-                var = md['var']
-                label = f'{var} [m]'
-            if 'var' in locals():
-                scatter = plt.scatter(df['x']/1000, df['y']/1000, c=df[var], cmap=cmaps, s=1)
-
-        x0, x1 = ax.get_xlim() if xlim == (None, None) else xlim
-        y0, y1 = ax.get_ylim() if ylim == (None, None) else ylim
-
-
-        x_extent = x1 - x0
-        y_extent = y1 - y0
-
-        aspect_ratio = y_extent / x_extent
-
-        # Set the figure size based on the aspect ratio and a scale factor
-        plt.gcf().set_size_inches(8 * scale_factor, 8 * aspect_ratio * scale_factor)
-        ax.set_xlabel('x [km]')
-        ax.set_ylabel('y [km]')
-        ax.set_xlim(*xlim)
-        ax.set_ylim(*ylim)
-        # self.scale_axes(ax, factor=1000)
-        plt.title(title, fontsize=30*scale_factor)
-
-        if 'scatter' in locals():
-            cbar = fig.colorbar(scatter, ax=ax, orientation='horizontal', pad=0.05, fraction=0.05)
-            cbar.ax.xaxis.set_ticks_position('bottom')
-            if 'label' in locals():
-                cbar.set_label(label)
-
-        plt.tight_layout()
-
-        with self.plot_context():
-            if save:
-                plt.savefig(save, dpi=200)
-                print('Figure saved as', save)
-            else:
-                plt.show()
 
 class MetadataResult:
     def __init__(self, metadata):

{anta_database-0.1.4 → anta_database-0.2.0}/anta_database/indexing/index_database.py

@@ -1,4 +1,5 @@
 import os
+from tqdm import tqdm
 import glob
 import pandas as pd
 import sqlite3
@@ -14,20 +15,13 @@ class IndexDatabase:
         ages = pd.read_csv(tab_file, header=None, sep='\t', names=['file', 'age'])
         return dict(zip(ages['file'], ages['age']))
 
-    def get_dict_vars(self, tab_file) -> dict:
-        ages = pd.read_csv(tab_file, header=None, sep='\t', names=['file', 'var'])
-        return dict(zip(ages['file'], ages['var']))
-
     def index_database(self):
         Authors_ages = {}
         for _, row in self.index.iterrows():
             ages = self.get_dict_ages(f"{self.db_dir}/{row.directory}/IRH_ages.tab")
             Authors_ages.update({f"{row.directory}": ages})
 
-        Extra_vars = {}
-        for _, row in self.index.iterrows():
-            vars = self.get_dict_ages(f"{self.db_dir}/{row.directory}/vars.tab")
-            Extra_vars.update({f"{row.directory}": vars})
+        var_list = pd.read_csv(f"{self.db_dir}/vars.csv").columns
 
         if os.path.exists(self.file_db):
             os.remove(self.file_db)
@@ -67,8 +61,8 @@ class IndexDatabase:
             )
         ''')
 
-        # Insert the metadata for each dataset into the table
-        for file in glob.glob(f'{self.db_dir}/**/**/*.pkl', recursive=True):
+        pkl_files = list(glob.glob(f'{self.db_dir}/**/**/*.pkl', recursive=True))
+        for file in tqdm(pkl_files, desc="Processing files"):
             dir_name, file_name = os.path.split(file)
             pkl_dir, trace_id = os.path.split(dir_name)
             author_dir, _ = os.path.split(pkl_dir)
@@ -76,7 +70,6 @@ class IndexDatabase:
             author = os.path.basename(author_dir)
             file_name_, ext = os.path.splitext(file_name)
             relative_file_path = f'{author}/pkl/{trace_id}/{file_name}'
-            print(relative_file_path)
 
             # Get the author's ID from the authors table
             cursor.execute('SELECT id FROM authors WHERE name = ?', (author,))
@@ -87,8 +80,8 @@ class IndexDatabase:
             else:
                 age = None
 
-            if file_name_ in Extra_vars[author]:
-                var = Extra_vars[author][file_name_]
+            if file_name_ in var_list:
+                var = file_name_
             else:
                 var = None
 

anta_database-0.2.0/anta_database/plotting/__init__.py

@@ -0,0 +1 @@
+from .plotting import *