Merge pull request #468 from jeromekelleher/finalise-dataset-import

Remove requirement for date_field

Author: jeromekelleher

.github/workflows/ci.yml

@@ -80,7 +80,7 @@ jobs:
            
       - name: Validate
         run: |
-          sc2ts validate -v testrun/dataset.zarr testrun/results/test/test_2020-02-02.ts
+          sc2ts validate -v --date-field=date testrun/dataset.zarr testrun/results/test/test_2020-02-02.ts
 
       - name: Info
         run: |

sc2ts/cli.py

@@ -161,14 +161,20 @@ def import_alignments(dataset, fastas, initialise, progress, verbose):
 @click.command()
 @click.argument("dataset", type=click.Path(dir_okay=True, file_okay=False))
 @click.argument("metadata", type=click.Path(dir_okay=False, file_okay=True))
+@click.option(
+    "--field-descriptions",
+    type=click.File(mode="r"),
+    default=None,
+    help="JSON formatted file of field descriptions",
+)
 @click.option(
     "--viridian",
     is_flag=True,
     help="Do some preprocessing appropriate for the Viridian metadata "
     "(Available at https://figshare.com/ndownloader/files/49694808)",
 )
 @verbose
-def import_metadata(dataset, metadata, viridian, verbose):
+def import_metadata(dataset, metadata, field_descriptions, viridian, verbose):
     """
     Import a CSV/TSV metadata file into the dataset.
     """
@@ -178,48 +184,14 @@ def import_metadata(dataset, metadata, viridian, verbose):
     if viridian:
         dtype = {"Artic_primer_version": str}
     df_in = pd.read_csv(metadata, sep="\t", dtype=dtype)
-    date_field = "date"
     index_field = "Run"
     if viridian:
         df_in = sc2ts.massage_viridian_metadata(df_in)
     df = df_in.set_index(index_field)
-    sc2ts.Dataset.add_metadata(dataset, df)
-
-
-@click.command()
-@click.argument("in_dataset", type=click.Path(dir_okay=True, file_okay=False))
-@click.argument("out_dataset", type=click.Path(dir_okay=True, file_okay=False))
-@click.option(
-    "--date-field", default="date", help="The metadata field to use for dates"
-)
-@click.option(
-    "-a",
-    "--additional-field",
-    default=[],
-    help="Additional fields to sort by",
-    multiple=True,
-)
-@chunk_cache_size
-@progress
-@verbose
-def reorder_dataset(
-    in_dataset,
-    out_dataset,
-    chunk_cache_size,
-    date_field,
-    additional_field,
-    progress,
-    verbose,
-):
-    """
-    Create a copy of the specified dataset where the samples are reordered by
-    date (and optionally other fields).
-    """
-    setup_logging(verbose)
-    ds = sc2ts.Dataset(
-        in_dataset, chunk_cache_size=chunk_cache_size, date_field=date_field
-    )
-    ds.reorder(out_dataset, show_progress=progress, additional_fields=additional_field)
+    d = {}
+    if field_descriptions is not None:
+        d = json.load(field_descriptions)
+    sc2ts.Dataset.add_metadata(dataset, df, field_descriptions=d)
 
 
 @click.command()
@@ -415,6 +387,11 @@ def infer(config_file, start, stop, force):
 @dataset
 @click.argument("ts_file")
 @deletions_as_missing
+@click.option(
+    "--date-field",
+    default=None,
+    help="Specify date field to use. Required for metadata.",
+)
 @click.option(
     "--genotypes/--no-genotypes",
     default=True,
@@ -440,6 +417,7 @@ def infer(config_file, start, stop, force):
 def validate(
     dataset,
     ts_file,
+    date_field,
     deletions_as_missing,
     genotypes,
     metadata,
@@ -453,7 +431,9 @@ def validate(
     setup_logging(verbose)
 
     ts = tszip.load(ts_file)
-    ds = sc2ts.Dataset(dataset, chunk_cache_size=chunk_cache_size)
+    ds = sc2ts.Dataset(
+        dataset, date_field=date_field, chunk_cache_size=chunk_cache_size
+    )
     if genotypes:
         sc2ts.validate_genotypes(ts, ds, deletions_as_missing, show_progress=True)
     if metadata:
@@ -564,7 +544,6 @@ def cli():
 
 cli.add_command(import_alignments)
 cli.add_command(import_metadata)
-cli.add_command(reorder_dataset)
 
 cli.add_command(info_dataset)
 cli.add_command(info_matches)

sc2ts/dataset.py

@@ -99,8 +99,6 @@ def __len__(self):
 class CachedMetadataMapping(collections.abc.Mapping):
     def __init__(self, root, sample_id_map, date_field, chunk_cache_size):
         self.sample_id_map = sample_id_map
-        self.sample_date = root[f"sample_{date_field}"][:].astype(str)
-        self.sample_date_array = root[f"sample_{date_field}"]
         self.sample_id = root["sample_id"][:].astype(str)
         self.sample_id_array = root["sample_id"]
         # Mapping of field name to Zarr array
@@ -114,6 +112,10 @@ def __init__(self, root, sample_id_map, date_field, chunk_cache_size):
         self.chunk_cache = {}
 
         logger.debug(f"Got {self.num_fields} metadata fields")
+        self.date_field = date_field
+        if date_field is not None:
+            self.sample_date = root[f"sample_{date_field}"][:].astype(str)
+            self.sample_date_array = root[f"sample_{date_field}"]
 
     @property
     def num_fields(self):
@@ -145,6 +147,8 @@ def get_metadata(self, j):
                 d[key] = bool(d[key])
             else:
                 d[key] = str(d[key])
+        if self.date_field is None:
+            raise ValueError("No date field set, cannot get metadata items")
         # For compatibility in the short term:
         d["date"] = self.sample_date[j]
         d["strain"] = self.sample_id[j]
@@ -178,7 +182,7 @@ class Variant:
 
 class Dataset(collections.abc.Mapping):
 
-    def __init__(self, path, chunk_cache_size=1, date_field="date", skip_metadata=False):
+    def __init__(self, path, chunk_cache_size=1, date_field=None):
         logger.info(f"Loading dateset @{path} using {date_field} as date field")
         self.date_field = date_field
         self.path = pathlib.Path(path)
@@ -196,13 +200,12 @@ def __init__(self, path, chunk_cache_size=1, date_field="date", skip_metadata=Fa
         self.haplotypes = CachedHaplotypeMapping(
             self.root, self.sample_id_map, chunk_cache_size
         )
-        if not skip_metadata:
-            self.metadata = CachedMetadataMapping(
-                self.root,
-                self.sample_id_map,
-                date_field,
-                chunk_cache_size=chunk_cache_size,
-            )
+        self.metadata = CachedMetadataMapping(
+            self.root,
+            self.sample_id_map,
+            date_field,
+            chunk_cache_size=chunk_cache_size,
+        )
 
     def __getitem__(self, key):
         return self.root[key]
@@ -432,7 +435,7 @@ def append_alignments(path, alignments):
         zarr.consolidate_metadata(store)
 
     @staticmethod
-    def add_metadata(path, df):
+    def add_metadata(path, df, field_descriptions=dict()):
         """
         Add metadata from the specified dataframe, indexed by sample ID.
         Each column will be added as a new array with prefix "sample_"
@@ -467,6 +470,8 @@ def add_metadata(path, df):
                 overwrite=True,
             )
             z.attrs["_ARRAY_DIMENSIONS"] = ["samples"]
+            z.attrs["description"] = field_descriptions.get(colname, "")
+
             z[:] = data
             logger.info(f"Wrote metadata array {z.name}")
 

tests/conftest.py

@@ -31,7 +31,6 @@ def fx_alignments_fasta(fx_data_cache):
 
 @pytest.fixture
 def fx_alignments_mafft_fasta(fx_data_cache):
-    # This is bgzipped so we can access directly
     cache_path = fx_data_cache / "alignments-mafft.fasta"
     if not cache_path.exists():
         with gzip.open("tests/data/alignments-mafft.fasta.gz") as src:
@@ -105,7 +104,7 @@ def fx_dataset(tmp_path, fx_data_cache, fx_alignments_fasta, fx_metadata_df):
         )
         sc2ts.Dataset.add_metadata(fs_path, fx_metadata_df)
         sc2ts.Dataset.create_zip(fs_path, cache_path)
-    return sc2ts.Dataset(cache_path)
+    return sc2ts.Dataset(cache_path, date_field="date")
 
 
 @pytest.fixture

tests/test_cli.py

@@ -52,12 +52,17 @@ def test_suite_data(self, tmp_path, fx_metadata_tsv, fx_alignments_fasta):
             catch_exceptions=False,
         )
         assert result.exit_code == 0
+        fields_path = tmp_path / "fields.json"
+        with open(fields_path, "w") as f:
+            f.write(json.dumps({"NO SUCH": "A", "Viridian_pangolin": "PANGO"}))
 
         result = runner.invoke(
             cli.cli,
-            f"import-metadata {ds_path} {fx_metadata_tsv} ",
+            f"import-metadata {ds_path} {fx_metadata_tsv} --field-descriptions={fields_path}",
             catch_exceptions=False,
         )
+        ds = sc2ts.Dataset(ds_path)
+        assert ds.metadata.fields["Viridian_pangolin"].attrs["description"] == "PANGO"
 
     def test_viridian_metadata(
         self, tmp_path, fx_raw_viridian_metadata_tsv, fx_alignments_fasta
@@ -379,7 +384,7 @@ def test_date(self, tmp_path, fx_ts_map, fx_dataset, date):
         runner = ct.CliRunner(mix_stderr=False)
         result = runner.invoke(
             cli.cli,
-            f"validate {fx_dataset.path} {ts_path} ",
+            f"validate {fx_dataset.path} {ts_path} --date-field=date",
             catch_exceptions=False,
         )
         assert result.exit_code == 0

tests/test_dataset.py

@@ -135,7 +135,10 @@ def test_add_metadata(self, tmp_path, fx_encoded_alignments, fx_metadata_df):
         path = tmp_path / "dataset.vcz"
         ds = sc2ts.Dataset.new(path)
         sc2ts.Dataset.append_alignments(path, fx_encoded_alignments)
-        sc2ts.Dataset.add_metadata(path, fx_metadata_df)
+        field_descriptions = {col: col.upper() for col in fx_metadata_df}
+        sc2ts.Dataset.add_metadata(
+            path, fx_metadata_df, field_descriptions=field_descriptions
+        )
 
         sg_ds = sgkit.load_dataset(path)
         assert dict(sg_ds.sizes) == {
@@ -147,7 +150,9 @@ def test_add_metadata(self, tmp_path, fx_encoded_alignments, fx_metadata_df):
         }
         df = fx_metadata_df.loc[sg_ds["sample_id"].values]
         for col in fx_metadata_df:
-            nt.assert_array_equal(df[col], sg_ds[f"sample_{col}"])
+            x = sg_ds[f"sample_{col}"]
+            nt.assert_array_equal(df[col], x)
+            assert x.attrs["description"] == field_descriptions[col]
 
     def test_create_zip(self, tmp_path, fx_encoded_alignments, fx_metadata_df):
 
@@ -283,7 +288,7 @@ def test_import(self, tmp_path, fx_encoded_alignments_mafft):
         path = tmp_path / "dataset.vcz"
         sc2ts.Dataset.new(path)
         sc2ts.Dataset.append_alignments(path, fx_encoded_alignments_mafft)
-        ds = sc2ts.Dataset(path, skip_metadata=True)
+        ds = sc2ts.Dataset(path)
         assert len(ds.haplotypes) == 19
         for k, v in fx_encoded_alignments_mafft.items():
             h = ds.haplotypes[k]
@@ -363,6 +368,12 @@ def test_known(self, fx_dataset):
         assert d["Genbank_N"] == -1
         assert d["Viridian_pangolin"] == "A"
 
+    def test_known_no_date_field(self, fx_dataset):
+        ds = sc2ts.Dataset(fx_dataset.path)
+
+        with pytest.raises(ValueError, match="No date field set"):
+            ds.metadata["SRR11772659"]
+
     @pytest.mark.parametrize(
         ["chunk_size", "cache_size"],
         [
@@ -382,7 +393,7 @@ def test_chunk_size_cache_size(
         sc2ts.Dataset.new(path, samples_chunk_size=chunk_size)
         sc2ts.Dataset.append_alignments(path, fx_encoded_alignments)
         sc2ts.Dataset.add_metadata(path, fx_metadata_df)
-        ds = sc2ts.Dataset(path, chunk_cache_size=cache_size)
+        ds = sc2ts.Dataset(path, chunk_cache_size=cache_size, date_field="date")
         for strain in fx_encoded_alignments.keys():
             row = fx_metadata_df.loc[strain]
             d1 = ds.metadata[strain]
@@ -406,6 +417,10 @@ def test_as_dataframe(self, fx_dataset, fx_metadata_df):
             data2 = df2[col]
             nt.assert_array_equal(data1.to_numpy(), data2.to_numpy())
 
+    def test_metadata_field_descriptions(self, fx_dataset):
+        for array in fx_dataset.metadata.fields.values():
+            assert array.attrs["description"] == ""
+
 
 class TestEncodeAlignment:
     @pytest.mark.parametrize(