Move track_data proto conversion functions to track_data_utils.

PiperOrigin-RevId: 796397700
Change-Id: I80043b8fa5ef330fc884372be5ddeb76725de0df

Author: padr31

src/alphagenome/data/track_data.py

@@ -15,18 +15,15 @@
 
 """Track data container analogous to AnnData."""
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Sequence
 import copy
 import dataclasses
 import enum
 from typing import Any, Union
 
-from alphagenome import tensor_utils
 from alphagenome import typing
 from alphagenome.data import genome
 from alphagenome.data import ontology
-from alphagenome.protos import dna_model_pb2
-from alphagenome.protos import tensor_pb2
 from jaxtyping import Bool, Float32, Int32  # pylint: disable=g-multiple-import, g-importing-member
 import numpy as np
 import pandas as pd
@@ -686,39 +683,6 @@ def __sub__(self, other: 'TrackData') -> 'TrackData':
         uns=self.uns,
     )
 
-  def to_protos(
-      self,
-      *,
-      bytes_per_chunk: int = 0,
-      compression_type: tensor_pb2.CompressionType = (
-          tensor_pb2.CompressionType.COMPRESSION_TYPE_NONE
-      ),
-  ) -> tuple[dna_model_pb2.TrackData, Sequence[tensor_pb2.TensorChunk]]:
-    """Serializes `TrackData` to protobuf messages.
-
-    Args:
-      bytes_per_chunk: The maximum number of bytes per tensor chunk.
-      compression_type: The compression type to use for the tensor chunks.
-
-    Returns:
-      A tuple containing the `TrackData` protobuf message and a sequence of
-      `TensorChunk` protobuf messages.
-    """
-    tensor, chunks = tensor_utils.pack_tensor(
-        self.values,
-        bytes_per_chunk=bytes_per_chunk,
-        compression_type=compression_type,
-    )
-    return (
-        dna_model_pb2.TrackData(
-            values=tensor,
-            metadata=metadata_to_proto(self.metadata).metadata,
-            resolution=self.resolution,
-            interval=self.interval.to_proto() if self.interval else None,
-        ),
-        chunks,
-    )
-
 
 def concat(
     track_datas: Sequence[TrackData],
@@ -869,199 +833,3 @@ def interleave(
       interval=track_datas[0].interval,
       uns={'num_interleaved_trackdatas': len(track_datas)},
   )
-
-
-def metadata_to_proto(
-    metadata: TrackMetadata,
-) -> dna_model_pb2.TracksMetadata:
-  """Converts track metadata to a `TracksMetadata` protobuf message.
-
-  Args:
-    metadata: A pandas DataFrame containing track metadata.
-
-  Returns:
-    A `TracksMetadata` protobuf message.
-  """
-  names = metadata['name']
-  default_values = [None] * len(names)
-
-  columns = zip(
-      metadata['name'],
-      metadata['strand'],
-      metadata.get('ontology_curie', default_values),
-      metadata.get('biosample_type', default_values),
-      metadata.get('biosample_name', default_values),
-      metadata.get('biosample_life_stage', default_values),
-      metadata.get('transcription_factor', default_values),
-      metadata.get('histone_mark', default_values),
-      metadata.get('gtex_tissue', default_values),
-      metadata.get('Assay title', default_values),
-      metadata.get('data_source', default_values),
-      metadata.get('genetically_modified', default_values),
-      metadata.get('endedness', default_values),
-      metadata.get('nonzero_mean', default_values),
-      strict=True,
-  )
-
-  metadata_protos = []
-  for (
-      name,
-      strand,
-      ontology_curie,
-      biosample_type,
-      biosample_name,
-      biosample_life_stage,
-      transcription_factor,
-      histone_mark,
-      gtex_tissue,
-      assay,
-      data_source,
-      genetically_modified,
-      endedness,
-      nonzero_mean,
-  ) in columns:
-    if biosample_type:
-      biosample = dna_model_pb2.Biosample(
-          name=biosample_name,
-          type=dna_model_pb2.BiosampleType.Value(
-              f'BIOSAMPLE_TYPE_{biosample_type.upper()}'
-          ),
-          stage=biosample_life_stage,
-      )
-    else:
-      biosample = None
-    if endedness is not None:
-      match endedness:
-        case 'paired':
-          endedness = dna_model_pb2.Endedness.ENDEDNESS_PAIRED
-        case 'single':
-          endedness = dna_model_pb2.Endedness.ENDEDNESS_SINGLE
-        case _:
-          raise ValueError(f'Unknown endedness: {endedness}')
-
-    metadata_protos.append(
-        dna_model_pb2.TrackMetadata(
-            name=name,
-            strand=genome.Strand.from_str(strand).to_proto()
-            if strand
-            else None,
-            ontology_term=ontology.from_curie(ontology_curie).to_proto()
-            if ontology_curie
-            else None,
-            biosample=biosample,
-            transcription_factor_code=transcription_factor
-            if isinstance(transcription_factor, str)
-            else None,
-            histone_mark_code=histone_mark
-            if isinstance(histone_mark, str)
-            else None,
-            gtex_tissue=gtex_tissue,
-            assay=assay,
-            data_source=data_source,
-            genetically_modified=genetically_modified,
-            endedness=endedness,
-            nonzero_mean=nonzero_mean,
-        )
-    )
-
-  return dna_model_pb2.TracksMetadata(metadata=metadata_protos)
-
-
-def metadata_from_proto(
-    proto: dna_model_pb2.TracksMetadata,
-) -> TrackMetadata:
-  """Creates track metadata from a `TracksMetadata` protobuf message.
-
-  Args:
-    proto: A `TracksMetadata` protobuf message.
-
-  Returns:
-    A pandas DataFrame containing track metadata.
-  """
-  metadata = []
-  for track_proto in proto.metadata:
-    track_metadata = {
-        'name': track_proto.name,
-        'strand': str(genome.Strand.from_proto(track_proto.strand)),
-    }
-
-    if track_proto.HasField('assay'):
-      track_metadata['Assay title'] = track_proto.assay
-
-    if track_proto.HasField('ontology_term'):
-      track_metadata['ontology_curie'] = ontology.from_proto(
-          track_proto.ontology_term
-      ).ontology_curie
-
-    if track_proto.HasField('biosample'):
-      track_metadata['biosample_name'] = track_proto.biosample.name
-      track_metadata['biosample_type'] = (
-          dna_model_pb2.BiosampleType.Name(track_proto.biosample.type)
-          .removeprefix('BIOSAMPLE_TYPE_')
-          .lower()
-      )
-      if track_proto.biosample.HasField('stage'):
-        track_metadata['biosample_life_stage'] = track_proto.biosample.stage
-
-    if track_proto.HasField('transcription_factor_code'):
-      track_metadata['transcription_factor'] = (
-          track_proto.transcription_factor_code
-      )
-
-    if track_proto.HasField('histone_mark_code'):
-      track_metadata['histone_mark'] = track_proto.histone_mark_code
-
-    if track_proto.HasField('gtex_tissue'):
-      track_metadata['gtex_tissue'] = track_proto.gtex_tissue
-
-    if track_proto.HasField('data_source'):
-      track_metadata['data_source'] = track_proto.data_source
-
-    if track_proto.HasField('endedness'):
-      track_metadata['endedness'] = (
-          dna_model_pb2.Endedness.Name(track_proto.endedness)
-          .removeprefix('ENDEDNESS_')
-          .lower()
-      )
-
-    if track_proto.HasField('genetically_modified'):
-      track_metadata['genetically_modified'] = track_proto.genetically_modified
-
-    if track_proto.HasField('nonzero_mean'):
-      track_metadata['nonzero_mean'] = track_proto.nonzero_mean
-
-    metadata.append(track_metadata)
-  if metadata:
-    return pd.DataFrame(metadata)
-  else:
-    return pd.DataFrame(columns=['name', 'strand'])
-
-
-def from_protos(
-    proto: dna_model_pb2.TrackData,
-    chunks: Iterable[tensor_pb2.TensorChunk] = (),
-    *,
-    interval: genome.Interval | None = None,
-) -> TrackData:
-  """Creates a `TrackData` object from protobuf messages.
-
-  Args:
-    proto: A `TrackData` protobuf message.
-    chunks: A sequence of `TensorChunk` protobuf messages.
-    interval: Optional `Interval` object representing the genomic region
-      containing the tracks. Only used if the proto does not have an interval.
-
-  Returns:
-    A `TrackData` object.
-  """
-  metadata = metadata_from_proto(
-      dna_model_pb2.TracksMetadata(metadata=proto.metadata)
-  )
-
-  values = tensor_utils.unpack_proto(proto.values, chunks)
-  values = tensor_utils.upcast_floating(values)
-  resolution = proto.resolution if proto.HasField('resolution') else 1
-  if proto.HasField('interval'):
-    interval = genome.Interval.from_proto(proto.interval)
-
-  return TrackData(values, metadata, resolution=resolution, interval=interval)