From 04aee293ab8081f530973148685e184be9016400 Mon Sep 17 00:00:00 2001
From: dilyong <dlo26@cam.ac.uk>
Date: Thu, 21 Mar 2024 01:02:25 +0000
Subject: [PATCH 1/2] Created a file unimpeded/unimpeded.py, which contains a
 class and a function called get for quick loading of datasets.

---
 unimpeded/unimpeded.py | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)
 create mode 100644 unimpeded/unimpeded.py

diff --git a/unimpeded/unimpeded.py b/unimpeded/unimpeded.py
new file mode 100644
index 0000000..d366a56
--- /dev/null
+++ b/unimpeded/unimpeded.py
@@ -0,0 +1,29 @@
+#%%
+from anesthetic import read_chains
+#%%
+class unimpeded:
+
+    # which __init__ is better?
+    """
+    def __init__(self, **kwargs):
+        self.model = kwargs.pop('model', None)
+        self.data = kwargs.pop('data', None)
+        self.method = kwargs.pop('method', None)
+    """
+    def __init__(self, model, dataset, method, location):
+        self.model = model
+        self.dataset = dataset
+        self.method = method
+        self.location = location
+    
+    def get(self):
+        if self.location == 'hpc':
+            samples = read_chains(f"/home/dlo26/rds/rds-dirac-dp192-63QXlf5HuFo/dlo26/{self.method}/{self.model}/{self.dataset}/{self.dataset}_polychord_raw/{self.dataset}") # for hpc
+        elif self.location == 'local':
+            samples = read_chains(f"../../{self.method}/{self.model}/{self.dataset}/{self.dataset}_polychord_raw/{self.dataset}")
+        return samples
+
+#%%
+database = unimpeded('klcdm', 'bao.sdss_dr16', 'ns', 'local')
+samples = database.get()
+# %%

From 293fa9c651e9509f209fbe3053420c59533b6c08 Mon Sep 17 00:00:00 2001
From: dilyong <dlo26@cam.ac.uk>
Date: Fri, 22 Mar 2024 15:33:49 +0000
Subject: [PATCH 2/2] Added tests/test_unimpeded.py. Updated
 unimpeded.unimpeded so now the class takes in a file path.

---
 tests/test_unimpeded.py | 17 ++++++++++++++
 unimpeded/unimpeded.py  | 49 ++++++++++++++++++++++-------------------
 2 files changed, 43 insertions(+), 23 deletions(-)
 create mode 100644 tests/test_unimpeded.py

diff --git a/tests/test_unimpeded.py b/tests/test_unimpeded.py
new file mode 100644
index 0000000..d3c88fc
--- /dev/null
+++ b/tests/test_unimpeded.py
@@ -0,0 +1,17 @@
+from unimpeded.unimpeded import unimpeded
+
+# These locations will be changed to Zenodo links in the future.
+loc = 'local'
+if loc == 'local':
+    location = '/Users/ongdily/Documents/Cambridge/project2/codes/'
+elif loc == 'hpc':
+    location = '/home/dlo26/rds/rds-dirac-dp192-63QXlf5HuFo/dlo26/'
+
+
+def test_get():
+    database = unimpeded(location=location)
+    samples = database.get(model='klcdm', dataset='bao.sdss_dr16', method='ns')
+    return samples
+
+
+samples = test_get()
diff --git a/unimpeded/unimpeded.py b/unimpeded/unimpeded.py
index d366a56..ca5fa38 100644
--- a/unimpeded/unimpeded.py
+++ b/unimpeded/unimpeded.py
@@ -1,29 +1,32 @@
-#%%
+"""Loading datasets."""
 from anesthetic import read_chains
-#%%
+
+
 class unimpeded:
+    """Loading datasets.
+
+    **kwargs
+    ----------
+    location : str
+        The location where the datasets are stored.
 
-    # which __init__ is better?
     """
     def __init__(self, **kwargs):
-        self.model = kwargs.pop('model', None)
-        self.data = kwargs.pop('data', None)
-        self.method = kwargs.pop('method', None)
-    """
-    def __init__(self, model, dataset, method, location):
-        self.model = model
-        self.dataset = dataset
-        self.method = method
-        self.location = location
-    
-    def get(self):
-        if self.location == 'hpc':
-            samples = read_chains(f"/home/dlo26/rds/rds-dirac-dp192-63QXlf5HuFo/dlo26/{self.method}/{self.model}/{self.dataset}/{self.dataset}_polychord_raw/{self.dataset}") # for hpc
-        elif self.location == 'local':
-            samples = read_chains(f"../../{self.method}/{self.model}/{self.dataset}/{self.dataset}_polychord_raw/{self.dataset}")
-        return samples
+        self.location = kwargs.pop('location', None)
+
+    def get(self, model, dataset, method):
+        """Get the samples from the dataset.
 
-#%%
-database = unimpeded('klcdm', 'bao.sdss_dr16', 'ns', 'local')
-samples = database.get()
-# %%
+        Parameters
+        ----------
+        method : str
+            The chosen method, ns or mcmc.
+        model : str
+            The chosen model, e.g. klcdm.
+        dataset : str
+            The chosen dataset, e.g. bao.sdss_dr16.
+        """
+        samples = read_chains(
+            self.location + f"{method}/{model}/{dataset}/"
+            f"{dataset}_polychord_raw/{dataset}")
+        return samples