Skip to content

Single cell memmap dataset

FileNames

Bases: str, Enum

Names of files that are generated in SingleCellCollection.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
42
43
44
45
46
47
48
49
50
51
52
53
54
class FileNames(str, Enum):
    """Names of files that are generated in SingleCellCollection."""

    DATA = "data.npy"
    COLPTR = "col_ptr.npy"
    ROWPTR = "row_ptr.npy"
    METADATA = "metadata.json"
    DTYPE = "dtypes.json"
    FEATURES = "features"
    VERSION = "version.json"
    NEIGHBOR_INDICES = "neighbor_indices.npy"
    NEIGHBOR_INDICES_PTR = "neighbor_indptr.npy"
    NEIGHBOR_VALUES = "neighbor_values.npy"

METADATA

Bases: str, Enum

Stored metadata.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
69
70
71
72
class METADATA(str, Enum):
    """Stored metadata."""

    NUM_ROWS = "num_rows"

Mode

Bases: str, Enum

Valid modes for the single cell memory mapped dataset.

The write append mode is 'w+' while the read append mode is 'r+'.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
57
58
59
60
61
62
63
64
65
66
class Mode(str, Enum):
    """Valid modes for the single cell memory mapped dataset.

    The write append mode is 'w+' while the read append mode is 'r+'.
    """

    CREATE_APPEND = "w+"
    READ_APPEND = "r+"
    READ = "r"
    CREATE = "w"

NeighborSamplingStrategy

Bases: str, Enum

Valid sampling strategies for neighbor selection.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
75
76
77
78
79
class NeighborSamplingStrategy(str, Enum):
    """Valid sampling strategies for neighbor selection."""

    RANDOM = "random"
    FIRST = "first"

SingleCellMemMapDataset

Bases: SingleCellRowDataset

Represents one or more AnnData matrices.

Data is stored in large, memory-mapped arrays that enables fast access of datasets larger than the available amount of RAM on a system. SCMMAP implements a consistent API defined in SingleCellRowDataset.

Attributes:

Name Type Description
data_path str

Location of np.memmap files to be loaded from or that will be

mode Mode

Whether the dataset will be read in (r+) from np.memmap files or

data Optional[ndarray]

A numpy array of the data

row_index Optional[ndarray]

A numpy array of row pointers

col_index Optional[ndarray]

A numpy array of column values

metadata Dict[str, int]

Various metadata about the dataset.

_feature_index RowFeatureIndex

The corresponding RowFeatureIndex where features are

dtypes Dict[FileNames, str]

A dictionary containing the datatypes of the data, row_index,

_version str

The version of the dataset

load_neighbors bool

Whether to load and utilize neighbor information from the 'neighbor_key' in AnnData's .obsp. Defaults to False.

neighbor_key str

The key in AnnData's .obsp containing the sparse adjacency matrix for neighbors. Defaults to 'next_cell_ids'.

neighbor_sampling_strategy str

Strategy for sampling neighbors ('random'). Defaults to 'random'.

fallback_to_identity bool

If a cell has no neighbors, whether to use the cell itself as its neighbor. Defaults to True.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
 182
 183
 184
 185
 186
 187
 188
 189
 190
 191
 192
 193
 194
 195
 196
 197
 198
 199
 200
 201
 202
 203
 204
 205
 206
 207
 208
 209
 210
 211
 212
 213
 214
 215
 216
 217
 218
 219
 220
 221
 222
 223
 224
 225
 226
 227
 228
 229
 230
 231
 232
 233
 234
 235
 236
 237
 238
 239
 240
 241
 242
 243
 244
 245
 246
 247
 248
 249
 250
 251
 252
 253
 254
 255
 256
 257
 258
 259
 260
 261
 262
 263
 264
 265
 266
 267
 268
 269
 270
 271
 272
 273
 274
 275
 276
 277
 278
 279
 280
 281
 282
 283
 284
 285
 286
 287
 288
 289
 290
 291
 292
 293
 294
 295
 296
 297
 298
 299
 300
 301
 302
 303
 304
 305
 306
 307
 308
 309
 310
 311
 312
 313
 314
 315
 316
 317
 318
 319
 320
 321
 322
 323
 324
 325
 326
 327
 328
 329
 330
 331
 332
 333
 334
 335
 336
 337
 338
 339
 340
 341
 342
 343
 344
 345
 346
 347
 348
 349
 350
 351
 352
 353
 354
 355
 356
 357
 358
 359
 360
 361
 362
 363
 364
 365
 366
 367
 368
 369
 370
 371
 372
 373
 374
 375
 376
 377
 378
 379
 380
 381
 382
 383
 384
 385
 386
 387
 388
 389
 390
 391
 392
 393
 394
 395
 396
 397
 398
 399
 400
 401
 402
 403
 404
 405
 406
 407
 408
 409
 410
 411
 412
 413
 414
 415
 416
 417
 418
 419
 420
 421
 422
 423
 424
 425
 426
 427
 428
 429
 430
 431
 432
 433
 434
 435
 436
 437
 438
 439
 440
 441
 442
 443
 444
 445
 446
 447
 448
 449
 450
 451
 452
 453
 454
 455
 456
 457
 458
 459
 460
 461
 462
 463
 464
 465
 466
 467
 468
 469
 470
 471
 472
 473
 474
 475
 476
 477
 478
 479
 480
 481
 482
 483
 484
 485
 486
 487
 488
 489
 490
 491
 492
 493
 494
 495
 496
 497
 498
 499
 500
 501
 502
 503
 504
 505
 506
 507
 508
 509
 510
 511
 512
 513
 514
 515
 516
 517
 518
 519
 520
 521
 522
 523
 524
 525
 526
 527
 528
 529
 530
 531
 532
 533
 534
 535
 536
 537
 538
 539
 540
 541
 542
 543
 544
 545
 546
 547
 548
 549
 550
 551
 552
 553
 554
 555
 556
 557
 558
 559
 560
 561
 562
 563
 564
 565
 566
 567
 568
 569
 570
 571
 572
 573
 574
 575
 576
 577
 578
 579
 580
 581
 582
 583
 584
 585
 586
 587
 588
 589
 590
 591
 592
 593
 594
 595
 596
 597
 598
 599
 600
 601
 602
 603
 604
 605
 606
 607
 608
 609
 610
 611
 612
 613
 614
 615
 616
 617
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637
 638
 639
 640
 641
 642
 643
 644
 645
 646
 647
 648
 649
 650
 651
 652
 653
 654
 655
 656
 657
 658
 659
 660
 661
 662
 663
 664
 665
 666
 667
 668
 669
 670
 671
 672
 673
 674
 675
 676
 677
 678
 679
 680
 681
 682
 683
 684
 685
 686
 687
 688
 689
 690
 691
 692
 693
 694
 695
 696
 697
 698
 699
 700
 701
 702
 703
 704
 705
 706
 707
 708
 709
 710
 711
 712
 713
 714
 715
 716
 717
 718
 719
 720
 721
 722
 723
 724
 725
 726
 727
 728
 729
 730
 731
 732
 733
 734
 735
 736
 737
 738
 739
 740
 741
 742
 743
 744
 745
 746
 747
 748
 749
 750
 751
 752
 753
 754
 755
 756
 757
 758
 759
 760
 761
 762
 763
 764
 765
 766
 767
 768
 769
 770
 771
 772
 773
 774
 775
 776
 777
 778
 779
 780
 781
 782
 783
 784
 785
 786
 787
 788
 789
 790
 791
 792
 793
 794
 795
 796
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
class SingleCellMemMapDataset(SingleCellRowDataset):
    """Represents one or more AnnData matrices.

    Data is stored in large, memory-mapped arrays that enables fast access of
    datasets larger than the available amount of RAM on a system. SCMMAP
    implements a consistent API defined in SingleCellRowDataset.

    Attributes:
        data_path: Location of np.memmap files to be loaded from or that will be
        created.
        mode: Whether the dataset will be read in (r+) from np.memmap files or
        written to np.memmap files (w+).
        data: A numpy array of the data
        row_index: A numpy array of row pointers
        col_index: A numpy array of column values
        metadata: Various metadata about the dataset.
        _feature_index: The corresponding RowFeatureIndex where features are
        stored
        dtypes: A dictionary containing the datatypes of the data, row_index,
        and col_index arrays.
        _version: The version of the dataset
        load_neighbors (bool, optional): Whether to load and utilize neighbor information
            from the 'neighbor_key' in AnnData's .obsp. Defaults to False.
        neighbor_key (str, optional): The key in AnnData's .obsp containing the
            sparse adjacency matrix for neighbors. Defaults to 'next_cell_ids'.
        neighbor_sampling_strategy (str, optional): Strategy for sampling neighbors ('random').
            Defaults to 'random'.
        fallback_to_identity (bool, optional): If a cell has no neighbors, whether
            to use the cell itself as its neighbor. Defaults to True.
    """

    def __init__(
        self,
        data_path: str,
        h5ad_path: Optional[str] = None,
        num_elements: Optional[int] = None,
        num_rows: Optional[int] = None,
        mode: Mode = Mode.READ_APPEND,
        paginated_load_cutoff: int = 10_000,
        load_block_row_size: int = 1_000_000,
        feature_index_name="feature_id",
        # --- Neighbor Args ---
        load_neighbors: bool = False,
        neighbor_key: str = "next_cell_ids",
        neighbor_sampling_strategy: str = NeighborSamplingStrategy.RANDOM,
        fallback_to_identity: bool = True,
    ) -> None:
        """Instantiate the class.

        Args:
            data_path: The location where the data np.memmap files are read from
            or stored.
            h5ad_path: Optional, the location of the h5_ad path.
            num_elements: The total number of elements in the array.
            num_rows: The number of rows in the data frame.
            mode: Whether to read or write from the data_path.
            paginated_load_cutoff: MB size on disk at which to load the h5ad structure with paginated load.
            load_block_row_size: Number of rows to load into memory with paginated load
            feature_index_name: The name of the features if the features are only stored in features_df.index.values
            # --- New Neighbor Args ---
            load_neighbors (bool, optional): Boolean to control to control whether to load and utilize neighbor information
            neighbor_key (str, optional): The key in AnnData's .obsp containing neighbor information.
            neighbor_sampling_strategy (str, optional): Sampling strategy for neighbors.
            fallback_to_identity (bool, optional): If a cell has no neighbors, whether to use the cell itself as its neighbor.
        """
        self._version: str = importlib.metadata.version("bionemo.scdl")
        self.data_path: str = data_path
        self.mode: Mode = mode
        self.paginated_load_cutoff = paginated_load_cutoff
        self.load_block_row_size = load_block_row_size
        self.feature_index_name = feature_index_name
        # Backing arrays
        self.data: Optional[np.ndarray] = None
        self.row_index: Optional[np.ndarray] = None
        self.row_index: Optional[np.ndarray] = None

        # Metadata and attributes
        self.metadata: Dict[str, int] = {}

        # Stores the Feature Index, which tracks
        # the original AnnData features (e.g., gene names)
        # and allows us to store ragged arrays in our SCMMAP structure.
        self._feature_index: RowFeatureIndex = RowFeatureIndex()

        # Variables for int packing / reduced precision
        self.dtypes: Dict[FileNames, str] = {
            f"{FileNames.DATA.value}": "float32",
            f"{FileNames.COLPTR.value}": "uint32",
            f"{FileNames.ROWPTR.value}": "uint64",
            f"{FileNames.NEIGHBOR_INDICES.value}": "uint32",
            f"{FileNames.NEIGHBOR_INDICES_PTR.value}": "uint64",
            f"{FileNames.NEIGHBOR_VALUES.value}": "float32",
        }

        # Neighbor configuration
        self.load_neighbors = load_neighbors
        self._has_neighbors = False
        if load_neighbors:
            self._init_neighbor_args(neighbor_key, neighbor_sampling_strategy, fallback_to_identity)

        if mode == Mode.CREATE_APPEND and os.path.exists(data_path):
            raise FileExistsError(f"Output directory already exists: {data_path}")

        if h5ad_path is not None and (data_path is not None and os.path.exists(data_path)):
            raise FileExistsError(
                "Invalid input; both an existing SCMMAP and an h5ad file were passed. "
                "Please pass either an existing SCMMAP or an h5ad file."
            )

        # If there is only a data path, and it exists already, load SCMMAP data.
        elif data_path is not None and os.path.exists(data_path):
            self.__init__obj()
            self.load(data_path)

        # If there is only an h5ad path, load the HDF5 data
        elif h5ad_path is not None:
            self.__init__obj()
            self.load_h5ad(h5ad_path)
        else:
            match num_rows, num_elements:
                case (int(), int()):
                    self.__init__obj()
                    self._init_arrs(num_elements=num_elements, num_rows=num_rows)
                case _:
                    raise ValueError("An np.memmap path, an h5ad path, or the number of elements and rows is required")

    def _init_neighbor_args(self, neighbor_key, neighbor_sampling_strategy, fallback_to_identity):
        # Neighbor tracking
        self._has_neighbors = False  # Track if neighbor data was successfully loaded/found

        self.neighbor_key = neighbor_key
        try:
            # Convert string to enum if a string was passed
            if isinstance(neighbor_sampling_strategy, str):
                neighbor_sampling_strategy = NeighborSamplingStrategy(neighbor_sampling_strategy)
            # Validate that it's a valid enum value
            if not isinstance(neighbor_sampling_strategy, NeighborSamplingStrategy):
                raise ValueError(f"Unsupported neighbor_sampling_strategy: {neighbor_sampling_strategy}")
        except ValueError:
            raise ValueError(f"Unsupported neighbor_sampling_strategy: {neighbor_sampling_strategy}")

        self.neighbor_sampling_strategy = neighbor_sampling_strategy
        self.fallback_to_identity = fallback_to_identity

    def __init__obj(self):
        """Initializes the data path and writes the version."""
        os.makedirs(self.data_path, exist_ok=True)

        # Write the version
        if not os.path.exists(f"{self.data_path}/{FileNames.VERSION.value}"):
            with open(f"{self.data_path}/{FileNames.VERSION.value}", "w") as vfi:
                json.dump(self.version(), vfi)

    def _init_arrs(self, num_elements: int, num_rows: int) -> None:
        self.mode = Mode.CREATE_APPEND
        data_arr, col_arr, row_arr = _create_compressed_sparse_row_memmaps(
            num_elements=num_elements,
            num_rows=num_rows,
            memmap_dir_path=Path(self.data_path),
            mode=self.mode,
            dtypes=self.dtypes,
        )
        self.data = data_arr
        self.col_index = col_arr
        self.row_index = row_arr

    def version(self) -> str:
        """Returns a version number.

        (following <major>.<minor>.<point> convention).
        """
        return self._version

    def _extract_neighbor_data(self, adata) -> bool:
        """Extracts neighbor data from AnnData.obsp object and saves to memmap files.

        Args:
            adata: AnnData object containing neighbor information
        Returns:
            bool: True if neighbor data was successfully loaded/found, False otherwise.
        """
        # Check if neighbor key exists in AnnData.obsp
        if self.neighbor_key not in adata.obsp:
            warnings.warn(f"Neighbor key '{self.neighbor_key}' not found in AnnData.obsp. Neighbor loading skipped.")
            return False

        logger.info(f"Extracting neighbor data from {self.neighbor_key} in AnnData.obsp")

        # Get the neighbor matrix from obsp
        neighbor_matrix = adata.obsp[self.neighbor_key]

        # Check if the neighbor matrix is a sparse matrix
        if not scipy.sparse.issparse(neighbor_matrix):
            raise ValueError(f"Neighbor matrix for key '{self.neighbor_key}' is not a sparse matrix.")

        # Initialize memory-mapped arrays for neighbor data with proper sizes
        indptr_len = len(neighbor_matrix.indptr)
        nnz = len(neighbor_matrix.indices)  # number of non-zero elements
        # No need to calculate data_len separately since it equals nnz

        # Create memory-mapped arrays for neighbor data
        self._neighbor_indptr = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_INDICES_PTR.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_INDICES_PTR.value}"],
            mode=Mode.CREATE_APPEND.value,
            shape=(indptr_len,),
        )

        self._neighbor_indices = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_INDICES.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_INDICES.value}"],
            mode=Mode.CREATE_APPEND.value,
            shape=(nnz,),
        )

        self._neighbor_data = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_VALUES.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_VALUES.value}"],
            mode=Mode.CREATE_APPEND.value,
            shape=(nnz,),
        )

        # Copy data into memory-mapped arrays (with dtype conversion)
        self._neighbor_indptr[:] = neighbor_matrix.indptr.astype(
            self.dtypes[f"{FileNames.NEIGHBOR_INDICES_PTR.value}"]
        )
        self._neighbor_indices[:] = neighbor_matrix.indices.astype(self.dtypes[f"{FileNames.NEIGHBOR_INDICES.value}"])
        self._neighbor_data[:] = neighbor_matrix.data.astype(self.dtypes[f"{FileNames.NEIGHBOR_VALUES.value}"])

        logger.info("Neighbor data extracted to memory-mapped arrays")
        return True

    def _extract_neighbor_data_paginated(self, adata) -> bool:
        """Extracts neighbor data using paginated approach for large datasets.

        Uses the same pattern as paginated_load_h5ad with binary file I/O and chunking
        to efficiently handle large neighbor matrices without loading everything at once.

        Args:
            adata: AnnData object containing neighbor information
        Returns:
            bool: True if neighbor data was successfully loaded/found, False otherwise.
        """
        # Check if neighbor key exists in AnnData.obsp
        if self.neighbor_key not in adata.obsp:
            warnings.warn(f"Neighbor key '{self.neighbor_key}' not found in AnnData.obsp. Neighbor loading skipped.")
            return False

        logger.info(f"Extracting neighbor data from {self.neighbor_key} in AnnData.obsp using chunked approach")

        # Get the neighbor matrix from obsp
        neighbor_matrix = adata.obsp[self.neighbor_key]

        # Check if the neighbor matrix is a sparse matrix
        if not scipy.sparse.issparse(neighbor_matrix):
            raise ValueError(f"Neighbor matrix for key '{self.neighbor_key}' is not a sparse matrix.")

        # First write indptr which gives us the structure - this is usually small enough to handle in one go
        memmap_dir_path = Path(self.data_path)
        with open(f"{memmap_dir_path}/{FileNames.NEIGHBOR_INDICES_PTR.value}", "wb") as indptr_file:
            # Convert to hardcoded dtype before writing
            indptr_converted = neighbor_matrix.indptr.astype(self.dtypes[f"{FileNames.NEIGHBOR_INDICES_PTR.value}"])
            indptr_file.write(indptr_converted.tobytes())

        # Get dimensions from indptr
        num_rows = len(neighbor_matrix.indptr) - 1
        # Process indices and data in chunks based on rows
        with (
            open(f"{memmap_dir_path}/{FileNames.NEIGHBOR_INDICES.value}", "wb") as indices_file,
            open(f"{memmap_dir_path}/{FileNames.NEIGHBOR_VALUES.value}", "wb") as data_file,
        ):
            for row_start in range(0, num_rows, self.load_block_row_size):
                row_end = min(row_start + self.load_block_row_size, num_rows)

                # Get slice of the matrix for this chunk of rows
                chunk = neighbor_matrix[row_start:row_end]

                # Convert to hardcoded dtypes before writing
                indices_converted = chunk.indices.astype(self.dtypes[f"{FileNames.NEIGHBOR_INDICES.value}"])
                data_converted = chunk.data.astype(self.dtypes[f"{FileNames.NEIGHBOR_VALUES.value}"])

                # Write chunk data to files
                indices_file.write(indices_converted.tobytes())
                data_file.write(data_converted.tobytes())

                logger.info(f"Processed neighbor data rows {row_start} to {row_end - 1}")

        # Then re-open as memory-mapped arrays with the final shapes
        self._neighbor_indptr = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_INDICES_PTR.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_INDICES_PTR.value}"],
            mode=Mode.READ_APPEND.value,
            shape=(len(neighbor_matrix.indptr),),
        )

        self._neighbor_indices = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_INDICES.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_INDICES.value}"],
            mode=Mode.READ_APPEND.value,
            shape=(len(neighbor_matrix.indices),),
        )

        self._neighbor_data = np.memmap(
            f"{self.data_path}/{FileNames.NEIGHBOR_VALUES.value}",
            dtype=self.dtypes[f"{FileNames.NEIGHBOR_VALUES.value}"],
            mode=Mode.READ_APPEND.value,
            shape=(len(neighbor_matrix.data),),
        )

        logger.info("Neighbor data extracted to memory-mapped arrays using chunked approach")
        return True

    def get_row(
        self,
        index: int,
        return_features: bool = False,
        feature_vars: Optional[List[str]] = None,
    ) -> Tuple[Tuple[np.ndarray, np.ndarray], List[np.ndarray]]:
        """Returns a given row in the dataset along with optional features.

        Args:
            index: The row to be returned. This is in the range of [0, num_rows)
            return_features: boolean that indicates whether to return features
            feature_vars: Optional, feature variables to extract
        Return:
            [Tuple[np.ndarray, np.ndarray]: data values and column pointes
            List[np.ndarray]: optional, corresponding features.
        """
        start = self.row_index[index]
        end = self.row_index[index + 1]
        values = self.data[start:end]
        columns = self.col_index[start:end]
        ret = (values, columns)
        if return_features:
            return ret, self._feature_index.lookup(index, select_features=feature_vars)[0]
        else:
            return ret, None

    def get_row_with_neighbor(
        self,
        index: int,
        return_features: bool = False,
        feature_vars: Optional[List[str]] = None,
    ) -> Dict[str, Union[Tuple[np.ndarray, np.ndarray], int, Optional[List[np.ndarray]]]]:
        """Returns a given row in the dataset along with optional features and neighbor data.

        Args:
            index: The row to be returned. This is in the range of [0, num_rows)
            return_features: Boolean that indicates whether to return features
            feature_vars: Optional, feature variables to extract

        Returns:
            Dict with keys:
            - 'current_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for current cell
            - 'next_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for neighbor cell
            - 'current_cell_index': int - Index of current cell
            - 'next_cell_index': int - Index of neighbor cell
            - 'features': List[np.ndarray] - Features if return_features is True, else None

        Raises:
            ValueError: If neighbor functionality is disabled or no neighbor data is available
        """
        # Validate neighbor availability since this function requires neighbors
        if not (self.load_neighbors and self._has_neighbors):
            raise ValueError(
                "Cannot include neighbor data: neighbor functionality is disabled or no neighbor data available"
            )

        # Get current cell data using the existing get_row function
        current_cell_data, features = self.get_row(index, return_features, feature_vars)

        # Sample neighbor and get its data
        neighbor_index = self.sample_neighbor_index(index)

        # Case where neighbor is the same as current cell
        if neighbor_index == index:
            next_cell_data = current_cell_data
        else:
            # Get neighbor cell data using the get_row function
            next_cell_data, _ = self.get_row(neighbor_index, False, None)

        # Return all data in a dictionary format
        return {
            "current_cell": current_cell_data,
            "next_cell": next_cell_data,
            "current_cell_index": index,
            "next_cell_index": neighbor_index,
            "features": features,
        }

    def get_row_padded(
        self,
        index: int,
        return_features: bool = False,
        feature_vars: Optional[List[str]] = None,
    ) -> Tuple[np.ndarray, List[np.ndarray]]:
        """Returns a padded version of a row in the dataset.

        A padded version is one where the a sparse array representation is
        converted to a conventional represenentation. Optionally, features are
        returned.

        Args:
            index: The row to be returned
            return_features: boolean that indicates whether to return features
            feature_vars: Optional, feature variables to extract
        Return:
            np.ndarray: conventional row representation
            List[np.ndarray]: optional, corresponding features.
        """
        (row_values, row_column_pointer), features = self.get_row(index, return_features, feature_vars)
        return (
            _pad_sparse_array(row_values, row_column_pointer, self._feature_index.number_vars_at_row(index)),
            features,
        )

    def get_row_padded_with_neighbor(
        self,
        index: int,
        return_features: bool = False,
        feature_vars: Optional[List[str]] = None,
    ) -> Dict[str, Union[np.ndarray, int, List[np.ndarray]]]:
        """Returns a padded version of a row with optional neighbor data.

        A padded version converts sparse representation to a dense array where
        missing values are filled with zeros.

        Args:
            index: The row to be returned
            return_features: Boolean that indicates whether to return features
            feature_vars: Optional, feature variables to extract

        Returns:
            Dict with keys:
            - 'current_cell': np.ndarray - Padded array for current cell
            - 'next_cell': np.ndarray - Padded array for neighbor cell
            - 'current_cell_index': int - Index of current cell
            - 'next_cell_index': int - Index of neighbor cell
            - 'features': List[np.ndarray] - Features if return_features is True, else None

        Raises:
            ValueError: If neighbor functionality is disabled or no neighbor data is available
        """
        # Validate neighbor availability since this function requires neighbors
        if not (self.load_neighbors and self._has_neighbors):
            raise ValueError(
                "Cannot include neighbor data: neighbor functionality is disabled or no neighbor data available"
            )

        # Get both current cell and neighbor data
        result = self.get_row_with_neighbor(index, return_features, feature_vars)

        # Get current cell padded array using get_row_padded
        curr_padded, _ = self.get_row_padded(index, False, None)

        # For neighbor, get the padded array
        next_idx = result["next_cell_index"]
        if next_idx == index:
            # If neighbor is the same as current cell, reuse the current padded array
            next_padded = curr_padded
        else:
            # Otherwise get the neighbor's padded array
            next_padded, _ = self.get_row_padded(next_idx, False, None)

        # Return in dictionary format
        return {
            "current_cell": curr_padded,
            "next_cell": next_padded,
            "current_cell_index": result["current_cell_index"],
            "next_cell_index": result["next_cell_index"],
            "features": result["features"],
        }

    def get_row_column(self, index: int, column: int, impute_missing_zeros: bool = True) -> Optional[float]:
        """Returns the value at a given index and the corresponding column.

        Args:
            index: The index to be returned
            column: The column to be returned
            impute_missing_zeros: boolean that indicates whether to set missing
            data to 0
        Return:
            A float that is the value in the array or None.
        """
        (row_values, row_column_pointer), _ = self.get_row(index)
        if column is not None:
            for col_index, col in enumerate(row_column_pointer):
                if col == column:
                    # return the value at this position
                    return row_values[col_index]
                elif col > column:
                    try:
                        raise ValueError(f"Column pointer {col} is larger than the column {column}.")
                    except ValueError:
                        break
            return 0.0 if impute_missing_zeros else None

    def features(self) -> Optional[RowFeatureIndex]:
        """Return the corresponding RowFeatureIndex."""
        return self._feature_index

    def _load_mmap_file_if_exists(self, file_path, dtype):
        if os.path.exists(file_path):
            return np.memmap(file_path, dtype=dtype, mode=self.mode)
        else:
            raise FileNotFoundError(f"The mmap file at {file_path} is missing")

    def load(self, stored_path: str) -> None:
        """Loads the data at store_path that is an np.memmap format.

        Args:
            stored_path: directory with np.memmap files
        Raises:
            FileNotFoundError if the corresponding directory or files are not
            found, or if the metadata file is not present.
        """
        if not os.path.exists(stored_path):
            raise FileNotFoundError(
                f"""Error: the specified data path to the mmap files {stored_path} does not exist.
                                    Specify an updated filepath or provide an h5ad path to the dataset. The data can
                                    be loaded with SingleCellMemMapDataset.load_h5ad. Alternatively, the class can be instantiated
                                    with  SingleCellMemMapDataset(<path to data that will be created>, h5ad_path=<path to h5ad file>"""
            )
        self.data_path = stored_path
        self.mode = Mode.READ_APPEND

        # Metadata is required, so we must check if it exists and fail if not.
        if not os.path.exists(f"{self.data_path}/{FileNames.METADATA.value}"):
            raise FileNotFoundError(
                f"Error: the metadata file {self.data_path}/{FileNames.METADATA.value} does not exist."
            )

        with open(f"{self.data_path}/{FileNames.METADATA.value}", Mode.READ_APPEND.value) as mfi:
            self.metadata = json.load(mfi)

        if os.path.exists(f"{self.data_path}/{FileNames.FEATURES.value}"):
            self._feature_index = RowFeatureIndex.load(f"{self.data_path}/{FileNames.FEATURES.value}")

        if os.path.exists(f"{self.data_path}/{FileNames.DTYPE.value}"):
            with open(f"{self.data_path}/{FileNames.DTYPE.value}") as dfi:
                self.dtypes = json.load(dfi)

        # mmap the existing arrays
        self.data = self._load_mmap_file_if_exists(
            f"{self.data_path}/{FileNames.DATA.value}", self.dtypes[f"{FileNames.DATA.value}"]
        )
        self.row_index = self._load_mmap_file_if_exists(
            f"{self.data_path}/{FileNames.ROWPTR.value}", dtype=self.dtypes[f"{FileNames.ROWPTR.value}"]
        )
        self.col_index = self._load_mmap_file_if_exists(
            f"{self.data_path}/{FileNames.COLPTR.value}", dtype=self.dtypes[f"{FileNames.COLPTR.value}"]
        )

        # Load neighbor data
        if self.load_neighbors:
            self._load_neighbor_memmaps()

    def _write_metadata(self) -> None:
        with open(f"{self.data_path}/{FileNames.METADATA.value}", f"{Mode.CREATE.value}") as mfi:
            json.dump(self.metadata, mfi)

    def regular_load_h5ad(
        self,
        anndata_path: str,
    ) -> Tuple[pd.DataFrame, int]:
        """Method for loading an h5ad file into memorySu and converting it to the SCDL format.

        Args:
            anndata_path: location of data to load
        Raises:
            NotImplementedError if the data is not in scipy.sparse.spmatrix format
            ValueError it there is not count data
        Returns:
            pd.DataFrame: var variables for features
            int: number of rows in the dataframe.

        """
        adata = ad.read_h5ad(anndata_path)  # slow

        # Check and load neighbor data
        # NOTE: More clear to have a check here and not call _extract_neighbor_data() if there no neighbors
        if self.load_neighbors:
            self._has_neighbors = self._extract_neighbor_data(adata)

        if not isinstance(adata.X, scipy.sparse.spmatrix):
            raise NotImplementedError("Error: dense matrix loading not yet implemented.")

        # Check if raw data is present
        raw = getattr(adata, "raw", None)
        count_data = None
        if raw is not None:
            # If it is, attempt to get the counts in the raw data.
            count_data = getattr(raw, "X", None)

        if count_data is None:
            # No raw counts were present, resort to normalized
            count_data = getattr(adata, "X")
        if count_data is None:
            raise ValueError("This file does not have count data")

        shape = count_data.shape
        num_rows = shape[0]

        num_elements_stored = count_data.nnz

        self.dtypes[f"{FileNames.DATA.value}"] = count_data.dtype

        # Create the arrays.
        self._init_arrs(num_elements_stored, num_rows)
        # Store data
        self.data[0:num_elements_stored] = count_data.data

        # Store the col idx array
        self.col_index[0:num_elements_stored] = count_data.indices.astype(int)

        # Store the row idx array
        self.row_index[0 : num_rows + 1] = count_data.indptr.astype(int)

        vars = adata.var
        adata.file.close()

        return vars, num_rows

    def paginated_load_h5ad(
        self,
        anndata_path: str,
    ) -> Tuple[pd.DataFrame, int]:
        """Method for block loading a larger h5ad file and converting it to the SCDL format.

        This should be used in the case when the entire anndata file cannot be loaded into memory.
        The anndata is loaded into memory load_block_row_size number of rows at a time. Each chunk
        is converted into numpy memory maps which are then concatenated together.

        Raises:
            NotImplementedError if the data is not loaded in the CSRDataset format.

        Returns:
            pd.DataFrame: var variables for features
            int: number of rows in the dataframe.
        """
        adata = ad.read_h5ad(anndata_path, backed=True)

        if self.load_neighbors:
            self._has_neighbors = self._extract_neighbor_data_paginated(adata)

        if not isinstance(adata.X, ad.experimental.CSRDataset):
            raise NotImplementedError("Non-sparse format cannot be loaded: {type(adata.X)}.")
        num_rows = adata.X.shape[0]

        self.dtypes[f"{FileNames.DATA.value}"] = adata.X.dtype

        # Read the row indices into a memory map.
        mode = Mode.CREATE_APPEND
        self.row_index = _create_row_memmaps(num_rows, Path(self.data_path), mode, self.dtypes)
        self.row_index[:] = adata.X._indptr.astype(int)

        # The data from each column and data chunk of the original anndata file is read in. This is saved into the final
        # location of the memmap file. In this step, it is saved in the binary file format.
        memmap_dir_path = Path(self.data_path)
        with (
            open(f"{memmap_dir_path}/{FileNames.COLPTR.value}", "wb") as col_file,
            open(f"{memmap_dir_path}/{FileNames.DATA.value}", "wb") as data_file,
        ):
            n_elements = 0
            for row_start in range(0, num_rows, self.load_block_row_size):
                # Write each array's data to the file in binary format
                col_block = adata.X[row_start : row_start + self.load_block_row_size].indices
                col_file.write(col_block.tobytes())

                data_block = adata.X[row_start : row_start + self.load_block_row_size].data
                data_file.write(data_block.tobytes())

                n_elements += len(data_block)

        # The column and data files are re-opened as memory-mapped arrays with the final shape
        mode = Mode.READ_APPEND
        self.col_index = np.memmap(
            f"{memmap_dir_path}/{FileNames.COLPTR.value}",
            self.dtypes[f"{FileNames.COLPTR.value}"],
            mode=mode,
            shape=(n_elements,),
        )
        self.data = np.memmap(
            f"{memmap_dir_path}/{FileNames.DATA.value}",
            dtype=self.dtypes[f"{FileNames.DATA.value}"],
            mode=mode,
            shape=(n_elements,),
        )
        vars = adata.var
        adata.file.close()

        return vars, num_rows

    def _load_neighbor_memmaps(self):
        try:
            # mmap the existing arrays
            self._neighbor_indices = self._load_mmap_file_if_exists(
                f"{self.data_path}/{FileNames.NEIGHBOR_INDICES.value}",
                self.dtypes[f"{FileNames.NEIGHBOR_INDICES.value}"],
            )
            self._neighbor_indptr = self._load_mmap_file_if_exists(
                f"{self.data_path}/{FileNames.NEIGHBOR_INDICES_PTR.value}",
                self.dtypes[f"{FileNames.NEIGHBOR_INDICES_PTR.value}"],
            )
            self._neighbor_data = self._load_mmap_file_if_exists(
                f"{self.data_path}/{FileNames.NEIGHBOR_VALUES.value}",
                self.dtypes[f"{FileNames.NEIGHBOR_VALUES.value}"],
            )

            self._has_neighbors = True

        except FileNotFoundError:
            # Neighbor files don't exist - this is OK if load_neighbors=False
            # or if dataset was created without neighbors
            self._has_neighbors = False
            if self.load_neighbors:
                warnings.warn("Neighbor loading was requested but neighbor files are missing")

    def load_h5ad(
        self,
        anndata_path: str,
    ) -> None:
        """Loads an existing AnnData archive from disk.

        This creates a new backing data structure which is saved.
        Note: the storage utilized will roughly double. Currently, the data must
        be in a scipy.sparse.spmatrix format.

        Args:
            anndata_path: location of data to load
        Raises:
            FileNotFoundError if the data path does not exist.
            NotImplementedError if the data is not in scipy.sparse.spmatrix
            format
            ValueError it there is not count data
        """
        if not os.path.exists(anndata_path):
            raise FileNotFoundError(f"Error: could not find h5ad path {anndata_path}")
        file_size_MB = os.path.getsize(anndata_path) / (1_024**2)

        if file_size_MB < self.paginated_load_cutoff:
            features_df, num_rows = self.regular_load_h5ad(anndata_path)
        else:
            features_df, num_rows = self.paginated_load_h5ad(anndata_path)
        if len(features_df.columns) > 0:
            features = {col: np.array(features_df[col].values) for col in features_df.columns}
        elif len(features_df.index) > 0:
            features = {self.feature_index_name: features_df.index.values}
        else:
            features = {}
        self._feature_index.append_features(n_obs=num_rows, features=features, label=anndata_path)
        self.save()

    def save(self, output_path: Optional[str] = None) -> None:
        """Saves the class to a given output path.

        Args:
            output_path: The location to save - not yet implemented and should
            be self.data_path

        Raises:
           NotImplementedError if output_path is not None.
        """
        if f"{METADATA.NUM_ROWS.value}" not in self.metadata:
            self.metadata[f"{METADATA.NUM_ROWS.value}"] = self.number_of_rows()

        self._write_metadata()
        # Write the feature index. This may not exist.
        self._feature_index.save(f"{self.data_path}/{FileNames.FEATURES.value}")

        # Ensure the object is in a valid state. These are saved at creation!
        for postfix in [
            f"{FileNames.VERSION.value}",
            f"{FileNames.DATA.value}",
            f"{FileNames.COLPTR.value}",
            f"{FileNames.ROWPTR.value}",
            f"{FileNames.FEATURES.value}",
        ]:
            if not os.path.exists(f"{self.data_path}/{postfix}"):
                raise FileNotFoundError(f"This file should exist from object creation: {self.data_path}/{postfix}")

        self.data.flush()  # NOTE: saves the data to disk, do the approach for neighbor data
        self.row_index.flush()
        self.col_index.flush()

        # Flush neighbor data to disk if it exists
        if self._has_neighbors and self._neighbor_indptr is not None:
            self._neighbor_indptr.flush()
            self._neighbor_indices.flush()
            self._neighbor_data.flush()

        if output_path is not None:
            raise NotImplementedError("Saving to separate path is not yet implemented.")

        return True

    def get_neighbor_indices_for_cell(self, cell_index: int) -> np.ndarray:
        """Returns the array of neighbor indices for a given cell.

        Args:
            cell_index: Index of the cell to get neighbors for

        Returns:
            np.ndarray: Array of neighbor indices, empty if no neighbors or neighbor data unavailable

        Raises:
            IndexError: If cell_index is out of bounds
            ValueError: If neighbor functionality was explicitly enabled but data is unavailable
        """
        if not (0 <= cell_index < self.number_of_rows()):
            raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

        # Check if neighbor functionality was requested but is unavailable
        if self.load_neighbors and not self._has_neighbors:
            raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

        if not self.load_neighbors or not self._has_neighbors or self._neighbor_indptr is None:
            return np.array([], dtype=int)  # Return empty array if neighbor data not available

        # Get neighbor indices using CSR format indptr and indices
        start = self._neighbor_indptr[cell_index]
        end = self._neighbor_indptr[cell_index + 1]
        return self._neighbor_indices[start:end]

    def get_neighbor_weights_for_cell(self, cell_index: int) -> np.ndarray:
        """Returns the array of neighbor weights (e.g., pseudotime differences) for a given cell.

        Args:
            cell_index: Index of the cell to get neighbor weights for

        Returns:
            np.ndarray: Array of weights corresponding to neighbors, empty if no neighbors

        Raises:
            IndexError: If cell_index is out of bounds
        """
        # Check if neighbor functionality was requested but is unavailable
        if self.load_neighbors and not self._has_neighbors:
            raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

        if (
            not self.load_neighbors
            or not self._has_neighbors
            or self._neighbor_indptr is None
            or self._neighbor_data is None
        ):
            return np.array([], dtype=float)

        if not (0 <= cell_index < self.number_of_rows()):
            raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

        # Get neighbor weights using CSR format indptr and data
        start = self._neighbor_indptr[cell_index]
        end = self._neighbor_indptr[cell_index + 1]
        return self._neighbor_data[start:end]

    def sample_neighbor_index(self, cell_index: int) -> int:
        """Samples a neighbor index for the given cell based on the configured sampling strategy.

        Args:
            cell_index: Index of the cell to sample a neighbor for

        Returns:
            int: Index of the sampled neighbor
                 If no neighbors exist and fallback_to_identity is True, returns cell_index

        Raises:
            ValueError: If an unsupported sampling strategy is specified
            IndexError: If cell_index is out of bounds
        """
        # Basic validation
        if not (0 <= cell_index < self.number_of_rows()):
            raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

        # Check if neighbor functionality was requested but is unavailable
        if self.load_neighbors and not self._has_neighbors:
            raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

        # Skip sampling if neighbor functionality is disabled
        if not self.load_neighbors:
            return cell_index  # Always return self as neighbor when neighbors disabled

        # Get the neighbor indices for this cell
        neighbor_indices = self.get_neighbor_indices_for_cell(cell_index)

        # If no neighbors found, handle according to fallback policy
        if len(neighbor_indices) == 0:
            if self.fallback_to_identity:
                return cell_index  # Return the cell itself
            else:
                # NOTE: implement fallback policy here if needed
                warnings.warn(
                    f"Cell {cell_index} has no neighbors and fallback_to_identity=False. "
                    f"Returning cell index itself anyway."
                )
                return cell_index  # Currently always return self if no neighbors

        # Sample neighbor based on strategy
        if self.neighbor_sampling_strategy == NeighborSamplingStrategy.RANDOM:
            # Simple random sampling with equal probability
            chosen_index = np.random.choice(neighbor_indices)
            return chosen_index
        elif self.neighbor_sampling_strategy == NeighborSamplingStrategy.FIRST:
            # First neighbor sampling
            return neighbor_indices[0]
        # NOTE: Future - Add weighted sampling strategy
        else:
            raise ValueError(f"Unsupported neighbor sampling strategy: {self.neighbor_sampling_strategy}")

    def get_neighbor_stats(self) -> dict:
        """Returns statistics about the neighbors in the dataset.

        Returns:
            dict: Dictionary with neighbor statistics:
                - has_neighbors: Whether dataset has neighbor data
                - total_connections: Total number of neighbor relationships
                - min_neighbors_per_cell: Minimum number of neighbors any cell has
                - max_neighbors_per_cell: Maximum number of neighbors any cell has
                - avg_neighbors_per_cell: Average number of neighbors per cell
                - cells_with_no_neighbors: Count of cells that have no neighbors
        """
        if not self._has_neighbors or self._neighbor_indptr is None or self._neighbor_indices is None:
            return {"has_neighbors": False}

        # Calculate stats based on CSR indptr (difference between consecutive elements)
        neighbor_counts = np.diff(self._neighbor_indptr)

        return {
            "has_neighbors": True,
            "total_connections": len(self._neighbor_indices),
            "min_neighbors_per_cell": int(np.min(neighbor_counts)),
            "max_neighbors_per_cell": int(np.max(neighbor_counts)),
            "avg_neighbors_per_cell": float(np.mean(neighbor_counts)),
            "cells_with_no_neighbors": int(np.sum(neighbor_counts == 0)),
        }

    def number_of_values(self) -> int:
        """Get the total number of values in the array.

        For each index, the length of the corresponding np.ndarray of features is counted.

        Returns:
            The sum of lengths of the features in every row
        """
        return sum(self._feature_index.number_of_values())

    def number_of_rows(self) -> int:
        """The number of rows in the dataset.

        Returns:
            The number of rows in the dataset
        Raises:
            ValueError if the length of the number of rows in the feature
            index does not correspond to the number of stored rows.
        """
        if len(self._feature_index) > 0 and self._feature_index.number_of_rows() != self.row_index.size - 1:
            raise ValueError(
                f"""The nuber of rows in the feature index {self._feature_index.number_of_rows()}
                             does not correspond to the number of rows in the row_index {self.row_index.size - 1}"""
            )
        return self._feature_index.number_of_rows()

    def number_nonzero_values(self) -> int:
        """Number of non zero entries in the dataset."""
        return self.data.size

    def __len__(self):
        """Return the number of rows."""
        return self.number_of_rows()

    def __getitem__(self, idx: int) -> torch.Tensor:
        """Get the row values located and index idx."""
        return torch.from_numpy(np.stack(self.get_row(idx)[0]))

    def number_of_variables(self) -> List[int]:
        """Get the number of features in every entry in the dataset.

        Returns:
            A list containing the lengths of the features in every row
        """
        feats = self._feature_index
        if len(feats) == 0:
            return [0]
        num_vars = feats.column_dims()
        return num_vars

    def shape(self) -> Tuple[int, List[int]]:
        """Get the shape of the dataset.

        This is the number of entries by the the length of the feature index
        corresponding to that variable.

        Returns:
            The number of elements in the dataset
            A list containing the number of variables for each row.
        """
        return self.number_of_rows(), self.number_of_variables()

    def concat(
        self,
        other_dataset: Union[list["SingleCellMemMapDataset"], "SingleCellMemMapDataset"],
        extend_copy_size: int = 10 * 1_024 * 1_024,
        output_path: str | None = None,
        destroy_on_copy: bool = False,
    ) -> None:
        """Concatenates one or a list of SingleCellMemMapDatasest to the existing one.

        The data is stored in the same place as for the original data set or at output_path
        if it is set. Then, at output_path or at self.data_path, there would be a saved
        SingleCellMemmpDataset, which can be read in with SingleCellMemmpDataset(output_path).

        Args:
            other_dataset: A SingleCellMemMapDataset or a list of
            SingleCellMemMapDatasets
            extend_copy_size: how much to copy in memory at once
            output_path: location to store new dataset
            destroy_on_copy: Whether to remove the current data_path

        Raises:
           ValueError if the other dataset(s) are not of the same version or
           something of another type is passed in.
        """
        # Verify the other dataset or datasets are of the same type.
        match other_dataset:
            case self.__class__():
                other_dataset = [other_dataset]
            case list():
                pass
            case _:
                raise ValueError(
                    f"Expecting either a {SingleCellMemMapDataset} or a list thereof. Actually got: {type(other_dataset)}"
                )

        for dataset in other_dataset:
            if self.version() != dataset.version():
                raise ValueError(
                    f"""Incompatable versions: input version: {dataset.version()},
            this version:  {self.version}"""
                )

        # Set our mode:
        self.mode: Mode = Mode.READ_APPEND
        if output_path is not None:
            if destroy_on_copy:
                shutil.move(self.data_path, output_path)
            else:
                shutil.copytree(self.data_path, output_path)
            self.data_path = output_path

        mmaps = []
        mmaps.extend(other_dataset)

        # Copy the data from self and other into the new arrays.
        cumulative_elements = self.number_nonzero_values()
        cumulative_rows = self.number_of_rows()
        for mmap in mmaps:
            destination_memmap = np.memmap(
                f"{mmap.data_path}/{FileNames.ROWPTR.value}_copy",
                dtype=self.dtypes[f"{FileNames.ROWPTR.value}"],
                mode="w+",
                shape=mmap.row_index.shape,
            )
            destination_memmap[:] = mmap.row_index[:]

            destination_memmap += int(cumulative_elements)

            destination_memmap.flush()
            if destroy_on_copy:
                os.remove(f"{mmap.data_path}/{FileNames.ROWPTR.value}")

            extend_files(
                f"{self.data_path}/{FileNames.ROWPTR.value}",
                f"{mmap.data_path}/{FileNames.ROWPTR.value}_copy",
                buffer_size_b=extend_copy_size,
                delete_file2_on_complete=True,
                offset=np.dtype(self.dtypes[f"{FileNames.ROWPTR.value}"]).itemsize,
            )

            extend_files(
                f"{self.data_path}/{FileNames.DATA.value}",
                f"{mmap.data_path}/{FileNames.DATA.value}",
                buffer_size_b=extend_copy_size,
                delete_file2_on_complete=destroy_on_copy,
            )
            extend_files(
                f"{self.data_path}/{FileNames.COLPTR.value}",
                f"{mmap.data_path}/{FileNames.COLPTR.value}",
                buffer_size_b=extend_copy_size,
                delete_file2_on_complete=destroy_on_copy,
            )
            self._feature_index.concat(mmap._feature_index)
            # Update counters
            cumulative_elements += mmap.number_nonzero_values()
            cumulative_rows += mmap.number_of_rows()

        # Reopen the data, colptr, and rowptr arrays
        self.data = np.memmap(
            f"{self.data_path}/{FileNames.DATA.value}",
            dtype=self.dtypes[f"{FileNames.DATA.value}"],
            shape=(cumulative_elements,),
            mode=Mode.READ_APPEND.value,
        )
        self.row_index = np.memmap(
            f"{self.data_path}/{FileNames.ROWPTR.value}",
            dtype=self.dtypes[f"{FileNames.ROWPTR.value}"],
            shape=(cumulative_rows + 1,),
            mode=Mode.READ_APPEND.value,
        )
        self.col_index = np.memmap(
            f"{self.data_path}/{FileNames.COLPTR.value}",
            dtype=self.dtypes[f"{FileNames.COLPTR.value}"],
            shape=(cumulative_elements,),
            mode=Mode.READ_APPEND.value,
        )
        self.save()

__getitem__(idx)

Get the row values located and index idx.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1152
1153
1154
def __getitem__(self, idx: int) -> torch.Tensor:
    """Get the row values located and index idx."""
    return torch.from_numpy(np.stack(self.get_row(idx)[0]))

__init__(data_path, h5ad_path=None, num_elements=None, num_rows=None, mode=Mode.READ_APPEND, paginated_load_cutoff=10000, load_block_row_size=1000000, feature_index_name='feature_id', load_neighbors=False, neighbor_key='next_cell_ids', neighbor_sampling_strategy=NeighborSamplingStrategy.RANDOM, fallback_to_identity=True)

Instantiate the class.

Parameters:

Name Type Description Default
data_path str

The location where the data np.memmap files are read from

required
h5ad_path Optional[str]

Optional, the location of the h5_ad path.

None
num_elements Optional[int]

The total number of elements in the array.

None
num_rows Optional[int]

The number of rows in the data frame.

None
mode Mode

Whether to read or write from the data_path.

READ_APPEND
paginated_load_cutoff int

MB size on disk at which to load the h5ad structure with paginated load.

10000
load_block_row_size int

Number of rows to load into memory with paginated load

1000000
feature_index_name

The name of the features if the features are only stored in features_df.index.values

'feature_id'
load_neighbors bool

Boolean to control to control whether to load and utilize neighbor information

False
neighbor_key str

The key in AnnData's .obsp containing neighbor information.

'next_cell_ids'
neighbor_sampling_strategy str

Sampling strategy for neighbors.

RANDOM
fallback_to_identity bool

If a cell has no neighbors, whether to use the cell itself as its neighbor.

True
Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
def __init__(
    self,
    data_path: str,
    h5ad_path: Optional[str] = None,
    num_elements: Optional[int] = None,
    num_rows: Optional[int] = None,
    mode: Mode = Mode.READ_APPEND,
    paginated_load_cutoff: int = 10_000,
    load_block_row_size: int = 1_000_000,
    feature_index_name="feature_id",
    # --- Neighbor Args ---
    load_neighbors: bool = False,
    neighbor_key: str = "next_cell_ids",
    neighbor_sampling_strategy: str = NeighborSamplingStrategy.RANDOM,
    fallback_to_identity: bool = True,
) -> None:
    """Instantiate the class.

    Args:
        data_path: The location where the data np.memmap files are read from
        or stored.
        h5ad_path: Optional, the location of the h5_ad path.
        num_elements: The total number of elements in the array.
        num_rows: The number of rows in the data frame.
        mode: Whether to read or write from the data_path.
        paginated_load_cutoff: MB size on disk at which to load the h5ad structure with paginated load.
        load_block_row_size: Number of rows to load into memory with paginated load
        feature_index_name: The name of the features if the features are only stored in features_df.index.values
        # --- New Neighbor Args ---
        load_neighbors (bool, optional): Boolean to control to control whether to load and utilize neighbor information
        neighbor_key (str, optional): The key in AnnData's .obsp containing neighbor information.
        neighbor_sampling_strategy (str, optional): Sampling strategy for neighbors.
        fallback_to_identity (bool, optional): If a cell has no neighbors, whether to use the cell itself as its neighbor.
    """
    self._version: str = importlib.metadata.version("bionemo.scdl")
    self.data_path: str = data_path
    self.mode: Mode = mode
    self.paginated_load_cutoff = paginated_load_cutoff
    self.load_block_row_size = load_block_row_size
    self.feature_index_name = feature_index_name
    # Backing arrays
    self.data: Optional[np.ndarray] = None
    self.row_index: Optional[np.ndarray] = None
    self.row_index: Optional[np.ndarray] = None

    # Metadata and attributes
    self.metadata: Dict[str, int] = {}

    # Stores the Feature Index, which tracks
    # the original AnnData features (e.g., gene names)
    # and allows us to store ragged arrays in our SCMMAP structure.
    self._feature_index: RowFeatureIndex = RowFeatureIndex()

    # Variables for int packing / reduced precision
    self.dtypes: Dict[FileNames, str] = {
        f"{FileNames.DATA.value}": "float32",
        f"{FileNames.COLPTR.value}": "uint32",
        f"{FileNames.ROWPTR.value}": "uint64",
        f"{FileNames.NEIGHBOR_INDICES.value}": "uint32",
        f"{FileNames.NEIGHBOR_INDICES_PTR.value}": "uint64",
        f"{FileNames.NEIGHBOR_VALUES.value}": "float32",
    }

    # Neighbor configuration
    self.load_neighbors = load_neighbors
    self._has_neighbors = False
    if load_neighbors:
        self._init_neighbor_args(neighbor_key, neighbor_sampling_strategy, fallback_to_identity)

    if mode == Mode.CREATE_APPEND and os.path.exists(data_path):
        raise FileExistsError(f"Output directory already exists: {data_path}")

    if h5ad_path is not None and (data_path is not None and os.path.exists(data_path)):
        raise FileExistsError(
            "Invalid input; both an existing SCMMAP and an h5ad file were passed. "
            "Please pass either an existing SCMMAP or an h5ad file."
        )

    # If there is only a data path, and it exists already, load SCMMAP data.
    elif data_path is not None and os.path.exists(data_path):
        self.__init__obj()
        self.load(data_path)

    # If there is only an h5ad path, load the HDF5 data
    elif h5ad_path is not None:
        self.__init__obj()
        self.load_h5ad(h5ad_path)
    else:
        match num_rows, num_elements:
            case (int(), int()):
                self.__init__obj()
                self._init_arrs(num_elements=num_elements, num_rows=num_rows)
            case _:
                raise ValueError("An np.memmap path, an h5ad path, or the number of elements and rows is required")

__init__obj()

Initializes the data path and writes the version.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
326
327
328
329
330
331
332
333
def __init__obj(self):
    """Initializes the data path and writes the version."""
    os.makedirs(self.data_path, exist_ok=True)

    # Write the version
    if not os.path.exists(f"{self.data_path}/{FileNames.VERSION.value}"):
        with open(f"{self.data_path}/{FileNames.VERSION.value}", "w") as vfi:
            json.dump(self.version(), vfi)

__len__()

Return the number of rows.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1148
1149
1150
def __len__(self):
    """Return the number of rows."""
    return self.number_of_rows()

concat(other_dataset, extend_copy_size=10 * 1024 * 1024, output_path=None, destroy_on_copy=False)

Concatenates one or a list of SingleCellMemMapDatasest to the existing one.

The data is stored in the same place as for the original data set or at output_path if it is set. Then, at output_path or at self.data_path, there would be a saved SingleCellMemmpDataset, which can be read in with SingleCellMemmpDataset(output_path).

Parameters:

Name Type Description Default
other_dataset Union[list[SingleCellMemMapDataset], SingleCellMemMapDataset]

A SingleCellMemMapDataset or a list of

required
extend_copy_size int

how much to copy in memory at once

10 * 1024 * 1024
output_path str | None

location to store new dataset

None
destroy_on_copy bool

Whether to remove the current data_path

False
Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
def concat(
    self,
    other_dataset: Union[list["SingleCellMemMapDataset"], "SingleCellMemMapDataset"],
    extend_copy_size: int = 10 * 1_024 * 1_024,
    output_path: str | None = None,
    destroy_on_copy: bool = False,
) -> None:
    """Concatenates one or a list of SingleCellMemMapDatasest to the existing one.

    The data is stored in the same place as for the original data set or at output_path
    if it is set. Then, at output_path or at self.data_path, there would be a saved
    SingleCellMemmpDataset, which can be read in with SingleCellMemmpDataset(output_path).

    Args:
        other_dataset: A SingleCellMemMapDataset or a list of
        SingleCellMemMapDatasets
        extend_copy_size: how much to copy in memory at once
        output_path: location to store new dataset
        destroy_on_copy: Whether to remove the current data_path

    Raises:
       ValueError if the other dataset(s) are not of the same version or
       something of another type is passed in.
    """
    # Verify the other dataset or datasets are of the same type.
    match other_dataset:
        case self.__class__():
            other_dataset = [other_dataset]
        case list():
            pass
        case _:
            raise ValueError(
                f"Expecting either a {SingleCellMemMapDataset} or a list thereof. Actually got: {type(other_dataset)}"
            )

    for dataset in other_dataset:
        if self.version() != dataset.version():
            raise ValueError(
                f"""Incompatable versions: input version: {dataset.version()},
        this version:  {self.version}"""
            )

    # Set our mode:
    self.mode: Mode = Mode.READ_APPEND
    if output_path is not None:
        if destroy_on_copy:
            shutil.move(self.data_path, output_path)
        else:
            shutil.copytree(self.data_path, output_path)
        self.data_path = output_path

    mmaps = []
    mmaps.extend(other_dataset)

    # Copy the data from self and other into the new arrays.
    cumulative_elements = self.number_nonzero_values()
    cumulative_rows = self.number_of_rows()
    for mmap in mmaps:
        destination_memmap = np.memmap(
            f"{mmap.data_path}/{FileNames.ROWPTR.value}_copy",
            dtype=self.dtypes[f"{FileNames.ROWPTR.value}"],
            mode="w+",
            shape=mmap.row_index.shape,
        )
        destination_memmap[:] = mmap.row_index[:]

        destination_memmap += int(cumulative_elements)

        destination_memmap.flush()
        if destroy_on_copy:
            os.remove(f"{mmap.data_path}/{FileNames.ROWPTR.value}")

        extend_files(
            f"{self.data_path}/{FileNames.ROWPTR.value}",
            f"{mmap.data_path}/{FileNames.ROWPTR.value}_copy",
            buffer_size_b=extend_copy_size,
            delete_file2_on_complete=True,
            offset=np.dtype(self.dtypes[f"{FileNames.ROWPTR.value}"]).itemsize,
        )

        extend_files(
            f"{self.data_path}/{FileNames.DATA.value}",
            f"{mmap.data_path}/{FileNames.DATA.value}",
            buffer_size_b=extend_copy_size,
            delete_file2_on_complete=destroy_on_copy,
        )
        extend_files(
            f"{self.data_path}/{FileNames.COLPTR.value}",
            f"{mmap.data_path}/{FileNames.COLPTR.value}",
            buffer_size_b=extend_copy_size,
            delete_file2_on_complete=destroy_on_copy,
        )
        self._feature_index.concat(mmap._feature_index)
        # Update counters
        cumulative_elements += mmap.number_nonzero_values()
        cumulative_rows += mmap.number_of_rows()

    # Reopen the data, colptr, and rowptr arrays
    self.data = np.memmap(
        f"{self.data_path}/{FileNames.DATA.value}",
        dtype=self.dtypes[f"{FileNames.DATA.value}"],
        shape=(cumulative_elements,),
        mode=Mode.READ_APPEND.value,
    )
    self.row_index = np.memmap(
        f"{self.data_path}/{FileNames.ROWPTR.value}",
        dtype=self.dtypes[f"{FileNames.ROWPTR.value}"],
        shape=(cumulative_rows + 1,),
        mode=Mode.READ_APPEND.value,
    )
    self.col_index = np.memmap(
        f"{self.data_path}/{FileNames.COLPTR.value}",
        dtype=self.dtypes[f"{FileNames.COLPTR.value}"],
        shape=(cumulative_elements,),
        mode=Mode.READ_APPEND.value,
    )
    self.save()

features()

Return the corresponding RowFeatureIndex.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
679
680
681
def features(self) -> Optional[RowFeatureIndex]:
    """Return the corresponding RowFeatureIndex."""
    return self._feature_index

get_neighbor_indices_for_cell(cell_index)

Returns the array of neighbor indices for a given cell.

Parameters:

Name Type Description Default
cell_index int

Index of the cell to get neighbors for

required

Returns:

Type Description
ndarray

np.ndarray: Array of neighbor indices, empty if no neighbors or neighbor data unavailable

Raises:

Type Description
IndexError

If cell_index is out of bounds

ValueError

If neighbor functionality was explicitly enabled but data is unavailable

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
def get_neighbor_indices_for_cell(self, cell_index: int) -> np.ndarray:
    """Returns the array of neighbor indices for a given cell.

    Args:
        cell_index: Index of the cell to get neighbors for

    Returns:
        np.ndarray: Array of neighbor indices, empty if no neighbors or neighbor data unavailable

    Raises:
        IndexError: If cell_index is out of bounds
        ValueError: If neighbor functionality was explicitly enabled but data is unavailable
    """
    if not (0 <= cell_index < self.number_of_rows()):
        raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

    # Check if neighbor functionality was requested but is unavailable
    if self.load_neighbors and not self._has_neighbors:
        raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

    if not self.load_neighbors or not self._has_neighbors or self._neighbor_indptr is None:
        return np.array([], dtype=int)  # Return empty array if neighbor data not available

    # Get neighbor indices using CSR format indptr and indices
    start = self._neighbor_indptr[cell_index]
    end = self._neighbor_indptr[cell_index + 1]
    return self._neighbor_indices[start:end]

get_neighbor_stats()

Returns statistics about the neighbors in the dataset.

Returns:

Name Type Description
dict dict

Dictionary with neighbor statistics: - has_neighbors: Whether dataset has neighbor data - total_connections: Total number of neighbor relationships - min_neighbors_per_cell: Minimum number of neighbors any cell has - max_neighbors_per_cell: Maximum number of neighbors any cell has - avg_neighbors_per_cell: Average number of neighbors per cell - cells_with_no_neighbors: Count of cells that have no neighbors

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
def get_neighbor_stats(self) -> dict:
    """Returns statistics about the neighbors in the dataset.

    Returns:
        dict: Dictionary with neighbor statistics:
            - has_neighbors: Whether dataset has neighbor data
            - total_connections: Total number of neighbor relationships
            - min_neighbors_per_cell: Minimum number of neighbors any cell has
            - max_neighbors_per_cell: Maximum number of neighbors any cell has
            - avg_neighbors_per_cell: Average number of neighbors per cell
            - cells_with_no_neighbors: Count of cells that have no neighbors
    """
    if not self._has_neighbors or self._neighbor_indptr is None or self._neighbor_indices is None:
        return {"has_neighbors": False}

    # Calculate stats based on CSR indptr (difference between consecutive elements)
    neighbor_counts = np.diff(self._neighbor_indptr)

    return {
        "has_neighbors": True,
        "total_connections": len(self._neighbor_indices),
        "min_neighbors_per_cell": int(np.min(neighbor_counts)),
        "max_neighbors_per_cell": int(np.max(neighbor_counts)),
        "avg_neighbors_per_cell": float(np.mean(neighbor_counts)),
        "cells_with_no_neighbors": int(np.sum(neighbor_counts == 0)),
    }

get_neighbor_weights_for_cell(cell_index)

Returns the array of neighbor weights (e.g., pseudotime differences) for a given cell.

Parameters:

Name Type Description Default
cell_index int

Index of the cell to get neighbor weights for

required

Returns:

Type Description
ndarray

np.ndarray: Array of weights corresponding to neighbors, empty if no neighbors

Raises:

Type Description
IndexError

If cell_index is out of bounds

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
def get_neighbor_weights_for_cell(self, cell_index: int) -> np.ndarray:
    """Returns the array of neighbor weights (e.g., pseudotime differences) for a given cell.

    Args:
        cell_index: Index of the cell to get neighbor weights for

    Returns:
        np.ndarray: Array of weights corresponding to neighbors, empty if no neighbors

    Raises:
        IndexError: If cell_index is out of bounds
    """
    # Check if neighbor functionality was requested but is unavailable
    if self.load_neighbors and not self._has_neighbors:
        raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

    if (
        not self.load_neighbors
        or not self._has_neighbors
        or self._neighbor_indptr is None
        or self._neighbor_data is None
    ):
        return np.array([], dtype=float)

    if not (0 <= cell_index < self.number_of_rows()):
        raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

    # Get neighbor weights using CSR format indptr and data
    start = self._neighbor_indptr[cell_index]
    end = self._neighbor_indptr[cell_index + 1]
    return self._neighbor_data[start:end]

get_row(index, return_features=False, feature_vars=None)

Returns a given row in the dataset along with optional features.

Parameters:

Name Type Description Default
index int

The row to be returned. This is in the range of [0, num_rows)

required
return_features bool

boolean that indicates whether to return features

False
feature_vars Optional[List[str]]

Optional, feature variables to extract

None

Return: [Tuple[np.ndarray, np.ndarray]: data values and column pointes List[np.ndarray]: optional, corresponding features.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
def get_row(
    self,
    index: int,
    return_features: bool = False,
    feature_vars: Optional[List[str]] = None,
) -> Tuple[Tuple[np.ndarray, np.ndarray], List[np.ndarray]]:
    """Returns a given row in the dataset along with optional features.

    Args:
        index: The row to be returned. This is in the range of [0, num_rows)
        return_features: boolean that indicates whether to return features
        feature_vars: Optional, feature variables to extract
    Return:
        [Tuple[np.ndarray, np.ndarray]: data values and column pointes
        List[np.ndarray]: optional, corresponding features.
    """
    start = self.row_index[index]
    end = self.row_index[index + 1]
    values = self.data[start:end]
    columns = self.col_index[start:end]
    ret = (values, columns)
    if return_features:
        return ret, self._feature_index.lookup(index, select_features=feature_vars)[0]
    else:
        return ret, None

get_row_column(index, column, impute_missing_zeros=True)

Returns the value at a given index and the corresponding column.

Parameters:

Name Type Description Default
index int

The index to be returned

required
column int

The column to be returned

required
impute_missing_zeros bool

boolean that indicates whether to set missing

True

Return: A float that is the value in the array or None.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
def get_row_column(self, index: int, column: int, impute_missing_zeros: bool = True) -> Optional[float]:
    """Returns the value at a given index and the corresponding column.

    Args:
        index: The index to be returned
        column: The column to be returned
        impute_missing_zeros: boolean that indicates whether to set missing
        data to 0
    Return:
        A float that is the value in the array or None.
    """
    (row_values, row_column_pointer), _ = self.get_row(index)
    if column is not None:
        for col_index, col in enumerate(row_column_pointer):
            if col == column:
                # return the value at this position
                return row_values[col_index]
            elif col > column:
                try:
                    raise ValueError(f"Column pointer {col} is larger than the column {column}.")
                except ValueError:
                    break
        return 0.0 if impute_missing_zeros else None

get_row_padded(index, return_features=False, feature_vars=None)

Returns a padded version of a row in the dataset.

A padded version is one where the a sparse array representation is converted to a conventional represenentation. Optionally, features are returned.

Parameters:

Name Type Description Default
index int

The row to be returned

required
return_features bool

boolean that indicates whether to return features

False
feature_vars Optional[List[str]]

Optional, feature variables to extract

None

Return: np.ndarray: conventional row representation List[np.ndarray]: optional, corresponding features.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
def get_row_padded(
    self,
    index: int,
    return_features: bool = False,
    feature_vars: Optional[List[str]] = None,
) -> Tuple[np.ndarray, List[np.ndarray]]:
    """Returns a padded version of a row in the dataset.

    A padded version is one where the a sparse array representation is
    converted to a conventional represenentation. Optionally, features are
    returned.

    Args:
        index: The row to be returned
        return_features: boolean that indicates whether to return features
        feature_vars: Optional, feature variables to extract
    Return:
        np.ndarray: conventional row representation
        List[np.ndarray]: optional, corresponding features.
    """
    (row_values, row_column_pointer), features = self.get_row(index, return_features, feature_vars)
    return (
        _pad_sparse_array(row_values, row_column_pointer, self._feature_index.number_vars_at_row(index)),
        features,
    )

get_row_padded_with_neighbor(index, return_features=False, feature_vars=None)

Returns a padded version of a row with optional neighbor data.

A padded version converts sparse representation to a dense array where missing values are filled with zeros.

Parameters:

Name Type Description Default
index int

The row to be returned

required
return_features bool

Boolean that indicates whether to return features

False
feature_vars Optional[List[str]]

Optional, feature variables to extract

None

Returns:

Type Description
Dict[str, Union[ndarray, int, List[ndarray]]]

Dict with keys:

Dict[str, Union[ndarray, int, List[ndarray]]]
  • 'current_cell': np.ndarray - Padded array for current cell
Dict[str, Union[ndarray, int, List[ndarray]]]
  • 'next_cell': np.ndarray - Padded array for neighbor cell
Dict[str, Union[ndarray, int, List[ndarray]]]
  • 'current_cell_index': int - Index of current cell
Dict[str, Union[ndarray, int, List[ndarray]]]
  • 'next_cell_index': int - Index of neighbor cell
Dict[str, Union[ndarray, int, List[ndarray]]]
  • 'features': List[np.ndarray] - Features if return_features is True, else None

Raises:

Type Description
ValueError

If neighbor functionality is disabled or no neighbor data is available

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
def get_row_padded_with_neighbor(
    self,
    index: int,
    return_features: bool = False,
    feature_vars: Optional[List[str]] = None,
) -> Dict[str, Union[np.ndarray, int, List[np.ndarray]]]:
    """Returns a padded version of a row with optional neighbor data.

    A padded version converts sparse representation to a dense array where
    missing values are filled with zeros.

    Args:
        index: The row to be returned
        return_features: Boolean that indicates whether to return features
        feature_vars: Optional, feature variables to extract

    Returns:
        Dict with keys:
        - 'current_cell': np.ndarray - Padded array for current cell
        - 'next_cell': np.ndarray - Padded array for neighbor cell
        - 'current_cell_index': int - Index of current cell
        - 'next_cell_index': int - Index of neighbor cell
        - 'features': List[np.ndarray] - Features if return_features is True, else None

    Raises:
        ValueError: If neighbor functionality is disabled or no neighbor data is available
    """
    # Validate neighbor availability since this function requires neighbors
    if not (self.load_neighbors and self._has_neighbors):
        raise ValueError(
            "Cannot include neighbor data: neighbor functionality is disabled or no neighbor data available"
        )

    # Get both current cell and neighbor data
    result = self.get_row_with_neighbor(index, return_features, feature_vars)

    # Get current cell padded array using get_row_padded
    curr_padded, _ = self.get_row_padded(index, False, None)

    # For neighbor, get the padded array
    next_idx = result["next_cell_index"]
    if next_idx == index:
        # If neighbor is the same as current cell, reuse the current padded array
        next_padded = curr_padded
    else:
        # Otherwise get the neighbor's padded array
        next_padded, _ = self.get_row_padded(next_idx, False, None)

    # Return in dictionary format
    return {
        "current_cell": curr_padded,
        "next_cell": next_padded,
        "current_cell_index": result["current_cell_index"],
        "next_cell_index": result["next_cell_index"],
        "features": result["features"],
    }

get_row_with_neighbor(index, return_features=False, feature_vars=None)

Returns a given row in the dataset along with optional features and neighbor data.

Parameters:

Name Type Description Default
index int

The row to be returned. This is in the range of [0, num_rows)

required
return_features bool

Boolean that indicates whether to return features

False
feature_vars Optional[List[str]]

Optional, feature variables to extract

None

Returns:

Type Description
Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]

Dict with keys:

Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]
  • 'current_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for current cell
Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]
  • 'next_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for neighbor cell
Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]
  • 'current_cell_index': int - Index of current cell
Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]
  • 'next_cell_index': int - Index of neighbor cell
Dict[str, Union[Tuple[ndarray, ndarray], int, Optional[List[ndarray]]]]
  • 'features': List[np.ndarray] - Features if return_features is True, else None

Raises:

Type Description
ValueError

If neighbor functionality is disabled or no neighbor data is available

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
def get_row_with_neighbor(
    self,
    index: int,
    return_features: bool = False,
    feature_vars: Optional[List[str]] = None,
) -> Dict[str, Union[Tuple[np.ndarray, np.ndarray], int, Optional[List[np.ndarray]]]]:
    """Returns a given row in the dataset along with optional features and neighbor data.

    Args:
        index: The row to be returned. This is in the range of [0, num_rows)
        return_features: Boolean that indicates whether to return features
        feature_vars: Optional, feature variables to extract

    Returns:
        Dict with keys:
        - 'current_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for current cell
        - 'next_cell': Tuple[np.ndarray, np.ndarray] - (values, columns) for neighbor cell
        - 'current_cell_index': int - Index of current cell
        - 'next_cell_index': int - Index of neighbor cell
        - 'features': List[np.ndarray] - Features if return_features is True, else None

    Raises:
        ValueError: If neighbor functionality is disabled or no neighbor data is available
    """
    # Validate neighbor availability since this function requires neighbors
    if not (self.load_neighbors and self._has_neighbors):
        raise ValueError(
            "Cannot include neighbor data: neighbor functionality is disabled or no neighbor data available"
        )

    # Get current cell data using the existing get_row function
    current_cell_data, features = self.get_row(index, return_features, feature_vars)

    # Sample neighbor and get its data
    neighbor_index = self.sample_neighbor_index(index)

    # Case where neighbor is the same as current cell
    if neighbor_index == index:
        next_cell_data = current_cell_data
    else:
        # Get neighbor cell data using the get_row function
        next_cell_data, _ = self.get_row(neighbor_index, False, None)

    # Return all data in a dictionary format
    return {
        "current_cell": current_cell_data,
        "next_cell": next_cell_data,
        "current_cell_index": index,
        "next_cell_index": neighbor_index,
        "features": features,
    }

load(stored_path)

Loads the data at store_path that is an np.memmap format.

Parameters:

Name Type Description Default
stored_path str

directory with np.memmap files

required

Raises: FileNotFoundError if the corresponding directory or files are not found, or if the metadata file is not present.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
def load(self, stored_path: str) -> None:
    """Loads the data at store_path that is an np.memmap format.

    Args:
        stored_path: directory with np.memmap files
    Raises:
        FileNotFoundError if the corresponding directory or files are not
        found, or if the metadata file is not present.
    """
    if not os.path.exists(stored_path):
        raise FileNotFoundError(
            f"""Error: the specified data path to the mmap files {stored_path} does not exist.
                                Specify an updated filepath or provide an h5ad path to the dataset. The data can
                                be loaded with SingleCellMemMapDataset.load_h5ad. Alternatively, the class can be instantiated
                                with  SingleCellMemMapDataset(<path to data that will be created>, h5ad_path=<path to h5ad file>"""
        )
    self.data_path = stored_path
    self.mode = Mode.READ_APPEND

    # Metadata is required, so we must check if it exists and fail if not.
    if not os.path.exists(f"{self.data_path}/{FileNames.METADATA.value}"):
        raise FileNotFoundError(
            f"Error: the metadata file {self.data_path}/{FileNames.METADATA.value} does not exist."
        )

    with open(f"{self.data_path}/{FileNames.METADATA.value}", Mode.READ_APPEND.value) as mfi:
        self.metadata = json.load(mfi)

    if os.path.exists(f"{self.data_path}/{FileNames.FEATURES.value}"):
        self._feature_index = RowFeatureIndex.load(f"{self.data_path}/{FileNames.FEATURES.value}")

    if os.path.exists(f"{self.data_path}/{FileNames.DTYPE.value}"):
        with open(f"{self.data_path}/{FileNames.DTYPE.value}") as dfi:
            self.dtypes = json.load(dfi)

    # mmap the existing arrays
    self.data = self._load_mmap_file_if_exists(
        f"{self.data_path}/{FileNames.DATA.value}", self.dtypes[f"{FileNames.DATA.value}"]
    )
    self.row_index = self._load_mmap_file_if_exists(
        f"{self.data_path}/{FileNames.ROWPTR.value}", dtype=self.dtypes[f"{FileNames.ROWPTR.value}"]
    )
    self.col_index = self._load_mmap_file_if_exists(
        f"{self.data_path}/{FileNames.COLPTR.value}", dtype=self.dtypes[f"{FileNames.COLPTR.value}"]
    )

    # Load neighbor data
    if self.load_neighbors:
        self._load_neighbor_memmaps()

load_h5ad(anndata_path)

Loads an existing AnnData archive from disk.

This creates a new backing data structure which is saved. Note: the storage utilized will roughly double. Currently, the data must be in a scipy.sparse.spmatrix format.

Parameters:

Name Type Description Default
anndata_path str

location of data to load

required

Raises: FileNotFoundError if the data path does not exist. NotImplementedError if the data is not in scipy.sparse.spmatrix format ValueError it there is not count data

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
def load_h5ad(
    self,
    anndata_path: str,
) -> None:
    """Loads an existing AnnData archive from disk.

    This creates a new backing data structure which is saved.
    Note: the storage utilized will roughly double. Currently, the data must
    be in a scipy.sparse.spmatrix format.

    Args:
        anndata_path: location of data to load
    Raises:
        FileNotFoundError if the data path does not exist.
        NotImplementedError if the data is not in scipy.sparse.spmatrix
        format
        ValueError it there is not count data
    """
    if not os.path.exists(anndata_path):
        raise FileNotFoundError(f"Error: could not find h5ad path {anndata_path}")
    file_size_MB = os.path.getsize(anndata_path) / (1_024**2)

    if file_size_MB < self.paginated_load_cutoff:
        features_df, num_rows = self.regular_load_h5ad(anndata_path)
    else:
        features_df, num_rows = self.paginated_load_h5ad(anndata_path)
    if len(features_df.columns) > 0:
        features = {col: np.array(features_df[col].values) for col in features_df.columns}
    elif len(features_df.index) > 0:
        features = {self.feature_index_name: features_df.index.values}
    else:
        features = {}
    self._feature_index.append_features(n_obs=num_rows, features=features, label=anndata_path)
    self.save()

number_nonzero_values()

Number of non zero entries in the dataset.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1144
1145
1146
def number_nonzero_values(self) -> int:
    """Number of non zero entries in the dataset."""
    return self.data.size

number_of_rows()

The number of rows in the dataset.

Returns:

Type Description
int

The number of rows in the dataset

Raises: ValueError if the length of the number of rows in the feature index does not correspond to the number of stored rows.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
def number_of_rows(self) -> int:
    """The number of rows in the dataset.

    Returns:
        The number of rows in the dataset
    Raises:
        ValueError if the length of the number of rows in the feature
        index does not correspond to the number of stored rows.
    """
    if len(self._feature_index) > 0 and self._feature_index.number_of_rows() != self.row_index.size - 1:
        raise ValueError(
            f"""The nuber of rows in the feature index {self._feature_index.number_of_rows()}
                         does not correspond to the number of rows in the row_index {self.row_index.size - 1}"""
        )
    return self._feature_index.number_of_rows()

number_of_values()

Get the total number of values in the array.

For each index, the length of the corresponding np.ndarray of features is counted.

Returns:

Type Description
int

The sum of lengths of the features in every row

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1118
1119
1120
1121
1122
1123
1124
1125
1126
def number_of_values(self) -> int:
    """Get the total number of values in the array.

    For each index, the length of the corresponding np.ndarray of features is counted.

    Returns:
        The sum of lengths of the features in every row
    """
    return sum(self._feature_index.number_of_values())

number_of_variables()

Get the number of features in every entry in the dataset.

Returns:

Type Description
List[int]

A list containing the lengths of the features in every row

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
def number_of_variables(self) -> List[int]:
    """Get the number of features in every entry in the dataset.

    Returns:
        A list containing the lengths of the features in every row
    """
    feats = self._feature_index
    if len(feats) == 0:
        return [0]
    num_vars = feats.column_dims()
    return num_vars

paginated_load_h5ad(anndata_path)

Method for block loading a larger h5ad file and converting it to the SCDL format.

This should be used in the case when the entire anndata file cannot be loaded into memory. The anndata is loaded into memory load_block_row_size number of rows at a time. Each chunk is converted into numpy memory maps which are then concatenated together.

Returns:

Name Type Description
DataFrame

pd.DataFrame: var variables for features

int int

number of rows in the dataframe.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
def paginated_load_h5ad(
    self,
    anndata_path: str,
) -> Tuple[pd.DataFrame, int]:
    """Method for block loading a larger h5ad file and converting it to the SCDL format.

    This should be used in the case when the entire anndata file cannot be loaded into memory.
    The anndata is loaded into memory load_block_row_size number of rows at a time. Each chunk
    is converted into numpy memory maps which are then concatenated together.

    Raises:
        NotImplementedError if the data is not loaded in the CSRDataset format.

    Returns:
        pd.DataFrame: var variables for features
        int: number of rows in the dataframe.
    """
    adata = ad.read_h5ad(anndata_path, backed=True)

    if self.load_neighbors:
        self._has_neighbors = self._extract_neighbor_data_paginated(adata)

    if not isinstance(adata.X, ad.experimental.CSRDataset):
        raise NotImplementedError("Non-sparse format cannot be loaded: {type(adata.X)}.")
    num_rows = adata.X.shape[0]

    self.dtypes[f"{FileNames.DATA.value}"] = adata.X.dtype

    # Read the row indices into a memory map.
    mode = Mode.CREATE_APPEND
    self.row_index = _create_row_memmaps(num_rows, Path(self.data_path), mode, self.dtypes)
    self.row_index[:] = adata.X._indptr.astype(int)

    # The data from each column and data chunk of the original anndata file is read in. This is saved into the final
    # location of the memmap file. In this step, it is saved in the binary file format.
    memmap_dir_path = Path(self.data_path)
    with (
        open(f"{memmap_dir_path}/{FileNames.COLPTR.value}", "wb") as col_file,
        open(f"{memmap_dir_path}/{FileNames.DATA.value}", "wb") as data_file,
    ):
        n_elements = 0
        for row_start in range(0, num_rows, self.load_block_row_size):
            # Write each array's data to the file in binary format
            col_block = adata.X[row_start : row_start + self.load_block_row_size].indices
            col_file.write(col_block.tobytes())

            data_block = adata.X[row_start : row_start + self.load_block_row_size].data
            data_file.write(data_block.tobytes())

            n_elements += len(data_block)

    # The column and data files are re-opened as memory-mapped arrays with the final shape
    mode = Mode.READ_APPEND
    self.col_index = np.memmap(
        f"{memmap_dir_path}/{FileNames.COLPTR.value}",
        self.dtypes[f"{FileNames.COLPTR.value}"],
        mode=mode,
        shape=(n_elements,),
    )
    self.data = np.memmap(
        f"{memmap_dir_path}/{FileNames.DATA.value}",
        dtype=self.dtypes[f"{FileNames.DATA.value}"],
        mode=mode,
        shape=(n_elements,),
    )
    vars = adata.var
    adata.file.close()

    return vars, num_rows

regular_load_h5ad(anndata_path)

Method for loading an h5ad file into memorySu and converting it to the SCDL format.

Parameters:

Name Type Description Default
anndata_path str

location of data to load

required

Raises: NotImplementedError if the data is not in scipy.sparse.spmatrix format ValueError it there is not count data Returns: pd.DataFrame: var variables for features int: number of rows in the dataframe.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
def regular_load_h5ad(
    self,
    anndata_path: str,
) -> Tuple[pd.DataFrame, int]:
    """Method for loading an h5ad file into memorySu and converting it to the SCDL format.

    Args:
        anndata_path: location of data to load
    Raises:
        NotImplementedError if the data is not in scipy.sparse.spmatrix format
        ValueError it there is not count data
    Returns:
        pd.DataFrame: var variables for features
        int: number of rows in the dataframe.

    """
    adata = ad.read_h5ad(anndata_path)  # slow

    # Check and load neighbor data
    # NOTE: More clear to have a check here and not call _extract_neighbor_data() if there no neighbors
    if self.load_neighbors:
        self._has_neighbors = self._extract_neighbor_data(adata)

    if not isinstance(adata.X, scipy.sparse.spmatrix):
        raise NotImplementedError("Error: dense matrix loading not yet implemented.")

    # Check if raw data is present
    raw = getattr(adata, "raw", None)
    count_data = None
    if raw is not None:
        # If it is, attempt to get the counts in the raw data.
        count_data = getattr(raw, "X", None)

    if count_data is None:
        # No raw counts were present, resort to normalized
        count_data = getattr(adata, "X")
    if count_data is None:
        raise ValueError("This file does not have count data")

    shape = count_data.shape
    num_rows = shape[0]

    num_elements_stored = count_data.nnz

    self.dtypes[f"{FileNames.DATA.value}"] = count_data.dtype

    # Create the arrays.
    self._init_arrs(num_elements_stored, num_rows)
    # Store data
    self.data[0:num_elements_stored] = count_data.data

    # Store the col idx array
    self.col_index[0:num_elements_stored] = count_data.indices.astype(int)

    # Store the row idx array
    self.row_index[0 : num_rows + 1] = count_data.indptr.astype(int)

    vars = adata.var
    adata.file.close()

    return vars, num_rows

sample_neighbor_index(cell_index)

Samples a neighbor index for the given cell based on the configured sampling strategy.

Parameters:

Name Type Description Default
cell_index int

Index of the cell to sample a neighbor for

required

Returns:

Name Type Description
int int

Index of the sampled neighbor If no neighbors exist and fallback_to_identity is True, returns cell_index

Raises:

Type Description
ValueError

If an unsupported sampling strategy is specified

IndexError

If cell_index is out of bounds

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
def sample_neighbor_index(self, cell_index: int) -> int:
    """Samples a neighbor index for the given cell based on the configured sampling strategy.

    Args:
        cell_index: Index of the cell to sample a neighbor for

    Returns:
        int: Index of the sampled neighbor
             If no neighbors exist and fallback_to_identity is True, returns cell_index

    Raises:
        ValueError: If an unsupported sampling strategy is specified
        IndexError: If cell_index is out of bounds
    """
    # Basic validation
    if not (0 <= cell_index < self.number_of_rows()):
        raise IndexError(f"Cell index {cell_index} out of bounds for dataset with {self.number_of_rows()} cells")

    # Check if neighbor functionality was requested but is unavailable
    if self.load_neighbors and not self._has_neighbors:
        raise ValueError("Neighbor functionality was enabled but no neighbor data is available")

    # Skip sampling if neighbor functionality is disabled
    if not self.load_neighbors:
        return cell_index  # Always return self as neighbor when neighbors disabled

    # Get the neighbor indices for this cell
    neighbor_indices = self.get_neighbor_indices_for_cell(cell_index)

    # If no neighbors found, handle according to fallback policy
    if len(neighbor_indices) == 0:
        if self.fallback_to_identity:
            return cell_index  # Return the cell itself
        else:
            # NOTE: implement fallback policy here if needed
            warnings.warn(
                f"Cell {cell_index} has no neighbors and fallback_to_identity=False. "
                f"Returning cell index itself anyway."
            )
            return cell_index  # Currently always return self if no neighbors

    # Sample neighbor based on strategy
    if self.neighbor_sampling_strategy == NeighborSamplingStrategy.RANDOM:
        # Simple random sampling with equal probability
        chosen_index = np.random.choice(neighbor_indices)
        return chosen_index
    elif self.neighbor_sampling_strategy == NeighborSamplingStrategy.FIRST:
        # First neighbor sampling
        return neighbor_indices[0]
    # NOTE: Future - Add weighted sampling strategy
    else:
        raise ValueError(f"Unsupported neighbor sampling strategy: {self.neighbor_sampling_strategy}")

save(output_path=None)

Saves the class to a given output path.

Parameters:

Name Type Description Default
output_path Optional[str]

The location to save - not yet implemented and should

None
Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
def save(self, output_path: Optional[str] = None) -> None:
    """Saves the class to a given output path.

    Args:
        output_path: The location to save - not yet implemented and should
        be self.data_path

    Raises:
       NotImplementedError if output_path is not None.
    """
    if f"{METADATA.NUM_ROWS.value}" not in self.metadata:
        self.metadata[f"{METADATA.NUM_ROWS.value}"] = self.number_of_rows()

    self._write_metadata()
    # Write the feature index. This may not exist.
    self._feature_index.save(f"{self.data_path}/{FileNames.FEATURES.value}")

    # Ensure the object is in a valid state. These are saved at creation!
    for postfix in [
        f"{FileNames.VERSION.value}",
        f"{FileNames.DATA.value}",
        f"{FileNames.COLPTR.value}",
        f"{FileNames.ROWPTR.value}",
        f"{FileNames.FEATURES.value}",
    ]:
        if not os.path.exists(f"{self.data_path}/{postfix}"):
            raise FileNotFoundError(f"This file should exist from object creation: {self.data_path}/{postfix}")

    self.data.flush()  # NOTE: saves the data to disk, do the approach for neighbor data
    self.row_index.flush()
    self.col_index.flush()

    # Flush neighbor data to disk if it exists
    if self._has_neighbors and self._neighbor_indptr is not None:
        self._neighbor_indptr.flush()
        self._neighbor_indices.flush()
        self._neighbor_data.flush()

    if output_path is not None:
        raise NotImplementedError("Saving to separate path is not yet implemented.")

    return True

shape()

Get the shape of the dataset.

This is the number of entries by the the length of the feature index corresponding to that variable.

Returns:

Type Description
int

The number of elements in the dataset

List[int]

A list containing the number of variables for each row.

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
def shape(self) -> Tuple[int, List[int]]:
    """Get the shape of the dataset.

    This is the number of entries by the the length of the feature index
    corresponding to that variable.

    Returns:
        The number of elements in the dataset
        A list containing the number of variables for each row.
    """
    return self.number_of_rows(), self.number_of_variables()

version()

Returns a version number.

(following .. convention).

Source code in bionemo/scdl/io/single_cell_memmap_dataset.py
348
349
350
351
352
353
def version(self) -> str:
    """Returns a version number.

    (following <major>.<minor>.<point> convention).
    """
    return self._version