Updates for DP2 (#71)

* Updates for DP2

* Fast fix

Author: hombit

src/uncle_val/datasets/dp1.py

@@ -130,7 +130,10 @@ def _read_ccd_visit_table(path, columns):
     if columns is None or "seeing" in columns:
         seeing = pa.array(df["seeing"])
         replace_value = np.nanmean(seeing)
-        df["seeing"] = pa.compute.replace_with_mask(seeing, pa.compute.is_nan(seeing), replace_value)
+        nan_seeing_mask = pa.compute.is_nan(seeing)
+        if isinstance(nan_seeing_mask, pa.ChunkedArray):
+            nan_seeing_mask = nan_seeing_mask.combine_chunks()
+        df["seeing"] = pa.compute.replace_with_mask(seeing, nan_seeing_mask, replace_value)
     if columns is None or "detectorId" in columns:
         detector_cols = _polar_encode_detector(df["detectorId"])
         df = df.assign(**detector_cols)

src/uncle_val/datasets/materialized.py

@@ -45,9 +45,10 @@ class MaterializedDataLoaderContext:
         It will be deleted on the context exit.
     """
 
-    def __init__(self, input_dataset: LSDBIterableDataset, tmp_dir: Path | str):
+    def __init__(self, input_dataset: LSDBIterableDataset, tmp_dir: Path | str, cleanup: bool):
         self.input_dataset = input_dataset
         self.tmp_dir = Path(tmp_dir)
+        self.cleanup = cleanup
 
     def _serialize_data(self):
         n_chunks = 0
@@ -74,4 +75,5 @@ def __enter__(self) -> DataLoader:
         )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        shutil.rmtree(self.tmp_dir)
+        if self.cleanup:
+            shutil.rmtree(self.tmp_dir)

src/uncle_val/pipelines/plotting.py

@@ -212,8 +212,11 @@ def _get_hists(
         device=torch.device(device),
     )
 
-    with Client(n_workers=n_workers, threads_per_worker=1, memory_limit="8GB") as client:
-        print(f"Dask Dashboard Link: {client.dashboard_link}")
+    with Client(n_workers=n_workers, threads_per_worker=1, memory_limit="64GB") as client:
+        try:
+            print(f"Dask Dashboard Link: {client.dashboard_link}")
+        except KeyError as e:
+            print(f"Cannot get Dask Dashboard Link: {e}")
         hists_df = hists.compute()
 
     return hists_df

src/uncle_val/pipelines/training_loop.py

@@ -129,7 +129,10 @@ def training_loop(
     model = model.to(device)
 
     with Client(n_workers=n_workers, memory_limit="8GB", threads_per_worker=1) as client:
-        print(f"Dask Dashboard Link: {client.dashboard_link}")
+        try:
+            print(f"Dask Dashboard Link: {client.dashboard_link}")
+        except KeyError as e:
+            print(f"Cannot get Dask Dashboard Link: {e}")
 
         validation_dataset_lsdb = LSDBIterableDataset(
             catalog=catalog,
@@ -144,7 +147,9 @@ def training_loop(
             device=device,
         )
 
-        with MaterializedDataLoaderContext(validation_dataset_lsdb, tmp_validation_dir) as val_dataloader:
+        with MaterializedDataLoaderContext(
+            validation_dataset_lsdb, tmp_validation_dir, cleanup=False
+        ) as val_dataloader:
             val_stats_future: Future | None = None
             mean_val_loss_i = 0