alethiotx.cellprofiler

 70def add_metadata(s3_path='.', csv_name='Cells.csv', columns=None, output_dir=None):
 71    """Add TIF metadata to a CellProfiler CSV.
 72
 73    Downloads ``csv_name`` and ``Image_enriched.csv`` from the S3 output
 74    directory (or reads from a local directory), joins them by
 75    ``FileName_*`` columns, and returns a pandas DataFrame.
 76
 77    ``ImageNumber`` cannot be used as a join key because CellProfiler
 78    parallel tasks each start numbering from 1, producing duplicates
 79    after ``merge_csv.py`` concatenates them.
 80
 81    :param s3_path: S3 URI or local path of the CellProfiler output
 82        directory (e.g. ``s3://example-bucket/my-experiment/`` or ``.``).
 83    :type s3_path: str
 84    :param csv_name: Name of the CSV to enrich
 85        (e.g. ``Cells.csv``, ``Nucleus.csv``, ``Cytoplasm.csv``)
 86    :type csv_name: str
 87    :param columns: Specific metadata columns to add from
 88        Image_enriched.csv.  If ``None``, all columns are added except
 89        ``ImageNumber``, ``FileName_*``, and ``PathName_*``.
 90    :type columns: list[str] or None
 91    :param output_dir: Directory to write ``<stem>_enriched.csv``.
 92        If ``None`` (the default), nothing is written to disk.
 93    :type output_dir: str or Path or None
 94    :return: The enriched data.
 95    :rtype: pandas.DataFrame
 96
 97    **Examples**
 98
 99    Return a DataFrame without writing to disk::
100
101        >>> from alethiotx.cellprofiler import add_metadata
102        >>> df = add_metadata("s3://example-bucket/my-experiment/", "Cells.csv")
103        >>> print(df.shape)
104        (1000, 250)
105
106    Add specific columns only::
107
108        >>> df = add_metadata(
109        ...     "s3://example-bucket/my-experiment/",
110        ...     "Nucleus.csv",
111        ...     columns=["PlateID", "Well", "Site", "Z_Step"],
112        ... )
113
114    Also write the merged CSV to disk::
115
116        >>> df = add_metadata(
117        ...     "s3://example-bucket/my-experiment/",
118        ...     "Cells.csv",
119        ...     output_dir="~/Downloads",
120        ... )
121    """
122    s3_base = s3_path.rstrip("/")
123    stem = Path(csv_name).stem
124
125    if output_dir is not None:
126        output_dir = Path(output_dir).expanduser()
127        output_path = output_dir / f"{stem}_enriched.csv"
128    else:
129        output_path = None
130
131    with tempfile.TemporaryDirectory() as tmpdir:
132        tmpdir = Path(tmpdir)
133        ie_local = tmpdir / "Image_enriched.csv"
134        csv_local = tmpdir / csv_name
135
136        # Download from S3 or resolve local paths
137        if s3_base.startswith("s3://"):
138            print(f"Downloading {s3_base}/Image_enriched.csv ...")
139            _s3_download(f"{s3_base}/Image_enriched.csv", ie_local)
140            print(f"Downloading {s3_base}/{csv_name} ...")
141            _s3_download(f"{s3_base}/{csv_name}", csv_local)
142        else:
143            local_base = Path(s3_base)
144            ie_local = local_base / "Image_enriched.csv"
145            csv_local = local_base / csv_name
146
147        # Read Image_enriched header
148        with open(ie_local) as f:
149            ie_header = csv.DictReader(f).fieldnames
150
151        # FileName columns are the join key (ImageNumber is unreliable
152        # because parallel CellProfiler tasks each start from 1)
153        fname_cols = [c for c in ie_header if c.startswith("FileName_")]
154        if not fname_cols:
155            raise ValueError(
156                "Image_enriched.csv must contain FileName_* columns for joining.\n"
157                f"Available columns: {ie_header}"
158            )
159
160        # Determine which metadata columns to add
161        if columns is not None:
162            missing = [c for c in columns if c not in ie_header]
163            if missing:
164                raise ValueError(
165                    f"Columns not found in Image_enriched.csv: {missing}\n"
166                    f"Use list_metadata_columns() to see available columns."
167                )
168            merge_cols = list(columns)
169        else:
170            # Exclude columns that already exist in the target CSV
171            exclude = {"ImageNumber"} | {
172                c for c in ie_header
173                if c.startswith(("FileName_", "PathName_"))
174            }
175            merge_cols = [c for c in ie_header if c not in exclude]
176
177        # Build lookup: filename tuple → metadata dict
178        lookup = {}
179        with open(ie_local) as f:
180            for row in csv.DictReader(f):
181                key = tuple(row.get(k, "") for k in fname_cols)
182                lookup[key] = {k: row[k] for k in merge_cols}
183
184        empty_meta = {k: "" for k in merge_cols}
185
186        # Stream CSV row-by-row: write to output (or temp) file
187        t0 = time.time()
188        n = 0
189
190        # Write to final destination or a temp file
191        write_path = output_path or Path(tmpdir) / f"{stem}_enriched.csv"
192
193        with open(csv_local) as fin, \
194             open(write_path, "w", newline="") as fout:
195            reader = csv.DictReader(fin)
196
197            # Verify target CSV has the same FileName columns
198            target_fname_cols = {c for c in reader.fieldnames
199                                 if c.startswith("FileName_")}
200            missing_fnames = set(fname_cols) - target_fname_cols
201            if missing_fnames:
202                raise ValueError(
203                    f"{csv_name} is missing FileName columns needed for "
204                    f"join: {sorted(missing_fnames)}"
205                )
206
207            writer = csv.DictWriter(
208                fout, fieldnames=reader.fieldnames + merge_cols,
209            )
210            writer.writeheader()
211
212            for row in reader:
213                key = tuple(row.get(k, "") for k in fname_cols)
214                row.update(lookup.get(key, empty_meta))
215                writer.writerow(row)
216                n += 1
217
218        elapsed = time.time() - t0
219        print(f"Merged {n:,} rows x {len(merge_cols)} metadata columns "
220              f"in {elapsed:.1f}s")
221
222        # Read back as DataFrame
223        df = pd.read_csv(write_path)
224
225    if output_path is not None:
226        size_mb = output_path.stat().st_size / 1024 / 1024
227        print(f"Saved to: {output_path} ({size_mb:.1f} MB)")
228
229    return df

35def list_metadata_columns(s3_path='.'):
36    """List available metadata columns in Image_enriched.csv.
37
38    Downloads only ``Image_enriched.csv`` from the S3 output directory
39    and returns the column names (excluding ``ImageNumber``).
40
41    :param s3_path: S3 URI of the CellProfiler output directory
42    :type s3_path: str
43    :return: List of available column names
44    :rtype: list[str]
45
46    **Example**
47
48    >>> from alethiotx.cellprofiler import list_metadata_columns
49    >>> cols = list_metadata_columns("s3://example-bucket/my-experiment/")
50    >>> print(cols[:5])
51    ['PlateID', 'Well', 'Site', 'Z_Step', 'z_position']
52    """
53    base = s3_path.rstrip("/")
54
55    # Support local directories as well as S3 URIs
56    if base.startswith("s3://"):
57        with tempfile.TemporaryDirectory() as tmpdir:
58            ie_local = Path(tmpdir) / "Image_enriched.csv"
59            _s3_download(f"{base}/Image_enriched.csv", ie_local)
60            with open(ie_local) as f:
61                header = csv.DictReader(f).fieldnames
62    else:
63        ie_local = Path(base) / "Image_enriched.csv"
64        with open(ie_local) as f:
65            header = csv.DictReader(f).fieldnames
66
67    return [c for c in header if c != "ImageNumber"]