diff --git a/adelie/data.py b/adelie/data.py index 8d44bf74..30bea72a 100644 --- a/adelie/data.py +++ b/adelie/data.py @@ -501,3 +501,148 @@ def snp_phased_ancestry( "group_sizes": group_sizes, "penalty": penalty, } + + +def snp_combine_r( + n: int, + s: int, + A: int, + *, + K: int =1, + glm: str ="gaussian", + sparsity: float =0.95, + one_ratio: float =0.25, + two_ratio: float =0.05, + zero_penalty: float =0, + snr: float =1, + seed: int =0, +): + """Creates a SNP unphased, ancestry dataset. + + - The groups and group sizes are generated randomly + such that ``G`` groups are created and the sum of the group sizes is ``p``. + - The calldata matrix ``X`` has sparsity ratio + ``1 - one_ratio - two_ratio`` + where ``one_ratio`` of the entries are randomly set to ``1`` + and ``two_ratio`` are randomly set to ``2``. + - The ancestry matrix randomly generates integers in the range ``[0, A)``. + - The true coefficients :math:`\\beta` is such that ``sparsity`` proportion of the entries are set to :math:`0`. + - The response ``y`` is generated from the GLM specified by ``glm``. + - The penalty factors are by default set to ``np.sqrt(group_sizes)``, + however if ``zero_penalty > 0``, a random set of penalties will be set to zero, + in which case, ``penalty`` is rescaled such that the :math:`\\ell_2` norm squared is ``p``. + + Parameters + ---------- + n : int + Number of data points. + s : int + Number of SNPs. + A : int + Number of ancestries. + K : int, optional + Number of classes for multi-response GLMs. + Default is ``1``. + glm : str, optional + GLM name. + It must be one of the following: + + - ``"binomial"`` + - ``"cox"`` + - ``"gaussian"`` + - ``"multigaussian"`` + - ``"multinomial"`` + - ``"poisson"`` + + Default is ``"gaussian"``. + sparsity : float, optional + Proportion of :math:`\\beta` entries to be zeroed out. + Default is ``0.95``. + one_ratio : float, optional + Proportion of the entries of ``X`` that is set to ``1``. + Default is ``0.25``. + two_ratio : float, optional + Proportion of the entries of ``X`` that is set to ``2``. + Default is ``0.05``. + zero_penalty : float, optional + Proportion of ``penalty`` entries to be zeroed out. + Default is ``0``. + snr : float, optional + Signal-to-noise ratio. + Default is ``1``. + seed : int, optional + Random seed. + Default is ``0``. + + Returns + ------- + data : dict + A dictionary containing the generated data: + + - ``"X"``: feature matrix. + - ``"ancestries"``: ancestry label of the same shape as ``X``. + - ``"y"``: response vector. + - ``"groups"``: mapping of group index to the starting column index of ``X``. + - ``"group_sizes"``: mapping of group index to the group size. + - ``"penalty"``: penalty factor for each group index. + """ + assert n >= 1 + assert s >= 1 + assert A >= 1 + assert sparsity >= 0 and sparsity <= 1 + assert one_ratio >= 0 and one_ratio <= 1 + assert two_ratio >= 0 and two_ratio <= 1 + assert zero_penalty >= 0 and zero_penalty <= 1 + assert snr > 0 + assert seed >= 0 + + np.random.seed(seed) + + nnz_ratio = one_ratio + two_ratio + one_ratio = one_ratio / nnz_ratio + two_ratio = two_ratio / nnz_ratio + nnz = int(nnz_ratio * n * s) + nnz_indices = np.random.choice(n * s, nnz, replace=False) + nnz_indices = np.random.permutation(nnz_indices) + one_indices = nnz_indices[:int(one_ratio * nnz)] + two_indices = nnz_indices[int(one_ratio * nnz):] + X = np.zeros((n, s), dtype=np.int8) + X.ravel()[one_indices] = 1 + X.ravel()[two_indices] = 2 + + # Generate ancestries for each haplotype at each SNP position + # This is completely independent of the mutations + ancestries = np.zeros((n, 2 * s), dtype=np.int8) + ancestries.ravel()[:] = np.random.choice(A, n * 2 * s, replace=True) + + # Each SNP has (1 + A) columns in the dense matrix: + # - 1 column for calldata + # - A columns for ancestry dosages + groups = (1 + A) * np.arange(s) + group_sizes = np.full(s, 1 + A, dtype=int) + + penalty = np.sqrt(group_sizes) + penalty[np.random.choice(s, int(zero_penalty * s), replace=False)] = 0 + penalty /= np.linalg.norm(penalty) / np.sqrt(s * (1 + A)) + + beta = np.random.normal(0, 1, (s, K)) + beta_nnz_indices = np.random.choice(s, int((1-sparsity) * s), replace=False) + X_sub = X[:, beta_nnz_indices] + beta_sub = beta[beta_nnz_indices] + + eta = X_sub @ beta_sub + glm = _sample_y( + glm=glm, + eta=eta, + beta=beta_sub, + snr=snr, + ) + + return { + "X": np.asfortranarray(X), + "ancestries": np.asfortranarray(ancestries), + "glm": glm, + "groups": groups, + "group_sizes": group_sizes, + "penalty": penalty, + } diff --git a/adelie/io.py b/adelie/io.py index 4849016a..1ec59072 100644 --- a/adelie/io.py +++ b/adelie/io.py @@ -192,3 +192,154 @@ def write( if error != "": raise RuntimeError(error) return total_bytes, benchmark + + +class snp_combine_r(core_io.IOSNPCombineR): + """IO handler for SNP unphased, ancestry matrix. + + A SNP unphased, ancestry matrix is a matrix + that combines unphased calldata and phased local ancestry information. + + Let :math:`X \\in \\mathbb{R}^{n \\times s(1+A)}` denote such a matrix + where + :math:`n` is the number of samples, + :math:`s` is the number of SNPs, + and :math:`A` is the number of ancestries. + Every :math:`1+A` (contiguous) columns is called a + *SNP block* corresponding to a single SNP + with the following structure: + - The first column contains the SNP data (0, 1, or 2) + - The next :math:`A` columns contain the unphased ancestry dosages for each ancestry (0, 1, or 2) + + For example, for SNP :math:`j`, the columns are: + - Column :math:`j(1+A)`: SNP data + - Columns :math:`j(1+A)+1` to :math:`j(1+A)+A`: Ancestry dosages + + Parameters + ---------- + filename : str + File name containing the SNP data in ``.snpdat`` format. + read_mode : str, optional + Reading mode of the SNP data. + It must be one of the following: + + - ``"file"``: reads the file using standard file IO. + This method is the most general and portable, + however, with large files, it is the slowest option. + - ``"mmap"``: reads the file using mmap. + This method is only supported on Linux and MacOS. + It is the most efficient way to read large files. + + Default is ``"file"``. + """ + def __init__( + self, + filename: str, + read_mode: str ="file", + ): + core_io.IOSNPCombineR.__init__(self, filename, read_mode) + + def write( + self, + calldata: np.ndarray, + ancestries: np.ndarray, + A: int, + n_threads: int =1, + *, + selected_ancestries: np.ndarray = None, + ): + """Writes a dense SNP unphased, ancestry matrix to the file in ``.snpdat`` format. + + .. note:: + The calldata and ancestries matrices must not contain + any missing values. + + Parameters + ---------- + calldata : (n, s) ndarray + SNP unphased calldata in dense format. + ``calldata[i, j]`` is the data for individual ``i`` and SNP ``j``. + It must only contain values in :math:`\\{0,1,2\\}`. + ancestries : (n, 2*s) ndarray + Local ancestry information in dense format. + ``ancestries[i, 2*j]`` and ``ancestries[i, 2*j+1]`` are the ancestries + for individual ``i`` and SNP ``j`` for each of the two haplotypes. + It must only contain values in :math:`\\{0,\\ldots, A-1\\}`. + A : int + Total number of ancestries present in ``ancestries`` (upper bound of labels). + n_threads : int, optional + Number of threads. + Default is ``1``. + selected_ancestries : ndarray of uint32, optional + If provided, restrict the ancestry dosage columns per SNP to these + ancestry labels (in the given order). The output block width per SNP + becomes ``1 + len(selected_ancestries)`` instead of ``1 + A``. + + Returns + ------- + total_bytes : int + Number of bytes written. + benchmark : dict + Dictionary of benchmark timings for each step of the serializer. + """ + if selected_ancestries is None: + ( + total_bytes, + benchmark, + error, + ) = core_io.IOSNPCombineR.write(self, calldata, ancestries, A, n_threads) + else: + selected_ancestries = np.asarray(selected_ancestries, dtype=np.uint32) + ( + total_bytes, + benchmark, + error, + ) = core_io.IOSNPCombineR.write(self, calldata, ancestries, A, selected_ancestries, n_threads) + if error != "": + raise RuntimeError(error) + return total_bytes, benchmark + + +class snp_combine_s(core_io.IOSNPCombineS): + """IO handler for SNP both-ancestry matrix. + + Combines: + - Mutated haplotype counts per ancestry (like phased ancestry summed across haps) + - Ancestry dosage counts per ancestry (like unphased ancestry dosages) + + Shape: (n, s * (2*A)). For each SNP j: + - Columns j*(2*A) .. j*(2*A)+(A-1): number of mutated haplotypes labeled with ancestry a + - Columns j*(2*A)+A .. j*(2*A)+(2*A-1): ancestry dosages per ancestry a + """ + def __init__( + self, + filename: str, + read_mode: str ="file", + ): + core_io.IOSNPCombineS.__init__(self, filename, read_mode) + + def write( + self, + calldata: np.ndarray, + ancestries: np.ndarray, + A: int, + n_threads: int =1, + *, + selected_ancestries: np.ndarray = None, + ): + if selected_ancestries is None: + ( + total_bytes, + benchmark, + error, + ) = core_io.IOSNPCombineS.write(self, calldata, ancestries, A, n_threads) + else: + selected_ancestries = np.asarray(selected_ancestries, dtype=np.uint32) + ( + total_bytes, + benchmark, + error, + ) = core_io.IOSNPCombineS.write(self, calldata, ancestries, A, selected_ancestries, n_threads) + if error != "": + raise RuntimeError(error) + return total_bytes, benchmark diff --git a/adelie/matrix.py b/adelie/matrix.py index 2d9324ae..c287b2f7 100644 --- a/adelie/matrix.py +++ b/adelie/matrix.py @@ -1298,6 +1298,93 @@ def __init__(self): return _snp_unphased() +def snp_combine_r( + io: io.snp_combine_r, + *, + n_threads: int =1, + dtype: Union[np.float32, np.float64] =np.float64, +): + """Creates a SNP unphased, ancestry matrix. + + The SNP unphased, ancestry matrix is a wrapper around + the corresponding IO handler :class:`adelie.io.snp_combine_r` + exposing some matrix operations. + + .. note:: + This matrix only works for naive method! + + Parameters + ---------- + io : snp_phased_ancestry + IO handler for SNP unphased, ancestry data. + n_threads : int, optional + Number of threads. + Default is ``1``. + dtype : Union[float32, float64], optional + Underlying value type. + Default is ``np.float64``. + + Returns + ------- + wrap + Wrapper matrix object. + + See Also + -------- + adelie.io.snp_combine_r + adelie.adelie_core.matrix.MatrixNaiveSNPCombineR32 + adelie.adelie_core.matrix.MatrixNaiveSNPCombineR64 + """ + dispatcher = { + np.float64: core.matrix.MatrixNaiveSNPCombineR64, + np.float32: core.matrix.MatrixNaiveSNPCombineR32, + } + core_base = dispatcher[dtype] + py_base = PyMatrixNaiveBase + + if not io.is_read: + io.read() + + class _snp_combine_r(core_base, py_base): + def __init__(self): + self._io = io + core_base.__init__(self, self._io, n_threads) + py_base.__init__(self, n_threads=n_threads) + + return _snp_combine_r() + + +def snp_combine_s( + io: io.snp_combine_s, + *, + n_threads: int =1, + dtype: Union[np.float32, np.float64] =np.float64, +): + """Creates a SNP both-ancestry matrix. + + For each SNP j, columns are divided into two blocks of A each: + - First A: mutated haplotype counts per ancestry + - Next A: ancestry dosage counts per ancestry + """ + dispatcher = { + np.float64: core.matrix.MatrixNaiveSNPCombineS64, + np.float32: core.matrix.MatrixNaiveSNPCombineS32, + } + core_base = dispatcher[dtype] + py_base = PyMatrixNaiveBase + + if not io.is_read: + io.read() + + class _snp_combine_s(core_base, py_base): + def __init__(self): + self._io = io + core_base.__init__(self, self._io, n_threads) + py_base.__init__(self, n_threads=n_threads) + + return _snp_combine_s() + + def sparse( mat: Union[csc_matrix, csr_matrix], *, diff --git a/adelie/src/include/adelie_core/io/io_snp_combine_r.hpp b/adelie/src/include/adelie_core/io/io_snp_combine_r.hpp new file mode 100644 index 00000000..c684e3b4 --- /dev/null +++ b/adelie/src/include/adelie_core/io/io_snp_combine_r.hpp @@ -0,0 +1,193 @@ +#pragma once +#include +#include + +#ifndef ADELIE_CORE_IO_SNP_COMBINE_R_TP +#define ADELIE_CORE_IO_SNP_COMBINE_R_TP \ + template +#endif +#ifndef ADELIE_CORE_IO_SNP_COMBINE_R +#define ADELIE_CORE_IO_SNP_COMBINE_R \ + IOSNPCombineR +#endif + +namespace adelie_core { +namespace io { + +template >> +class IOSNPCombineR : public IOSNPBase +{ +public: + using base_t = IOSNPBase; + using outer_t = uint64_t; + using inner_t = uint32_t; + using chunk_inner_t = uint8_t; + using value_t = int8_t; + using vec_outer_t = util::rowvec_type; + using vec_inner_t = util::rowvec_type; + using vec_value_t = util::rowvec_type; + using rowarr_value_t = util::rowarr_type; + using colarr_value_t = util::colarr_type; + using typename base_t::bool_t; + using typename base_t::buffer_t; + + // static constexpr size_t n_bits_per_byte = 8; + // static constexpr size_t chunk_size = ( + // // casting helps MSVC with warning C4293 + // static_cast(1UL) << (n_bits_per_byte * sizeof(chunk_inner_t)) + // ); + // NOTE: Each chunk stores at most one *byte* per allele occurrence + // (dosage value 1) within a 128-sample window. In the worst case, where + // all samples in the window have dosage == 2, we emit 256 bytes which + // still fits in an unsigned 8-bit counter (encoded as *nnz − 1*). + // Reducing the chunk size from 256→128 therefore guarantees that the + // counter cannot overflow and avoids downstream buffer overruns that + // manifested as segmentation faults for large cohorts (n≈10 k). + static constexpr size_t chunk_size = 128; + static constexpr size_t n_bits_per_byte = 8; // retained for compatibility + +protected: + static constexpr size_t _max_inner = ( + // casting helps MSVC with warning C4293 + static_cast(1UL) << (n_bits_per_byte * sizeof(inner_t)) + ); + + using base_t::throw_no_read; + using base_t::fopen_safe; + using base_t::is_big_endian; + using base_t::_buffer; + using base_t::_filename; + using base_t::_is_read; + + outer_t _rows; + outer_t _snps; + outer_t _ancestries; + outer_t _cols; + vec_outer_t _nnz; + vec_outer_t _outer; + +public: + using iterator = IOSNPChunkIterator< + chunk_size, inner_t, chunk_inner_t + >; + + // + // Allow random-access "calls" so utils.hpp can do io(i,snp,anc) + // + inline value_t operator()(size_t sample, + size_t snp, + size_t anc) const + { + // scan the non-zero entries for (snp,anc) + for (auto it = begin(snp, anc), e = end(snp, anc); it != e; ++it) { + if (static_cast(*it) == sample) { + return value_t(1); + } + } + return value_t(0); + } + + using base_t::base_t; + + size_t read() override; + + outer_t rows() const { + if (!_is_read) throw_no_read(); + return _rows; + } + + outer_t snps() const { + if (!_is_read) throw_no_read(); + return _snps; + } + + outer_t ancestries() const + { + if (!_is_read) throw_no_read(); + return _ancestries; + } + + outer_t cols() const + { + if (!_is_read) throw_no_read(); + return _cols; + } + + Eigen::Ref nnz() const + { + if (!_is_read) throw_no_read(); + return _nnz; + } + + Eigen::Ref outer() const + { + if (!_is_read) throw_no_read(); + return _outer; + } + + Eigen::Ref col(int j) const + { + return Eigen::Map( + _buffer.data() + _outer[j], + _outer[j+1] - _outer[j] + ); + } + + const char* col_anc(int j, int anc) const + { + const auto _col = col(j); + return ( + _col.data() + + internal::read_as(_col.data() + sizeof(outer_t) * anc) + ); + } + + inner_t n_chunks(int j, int anc) const + { + const auto* _col_anc = col_anc(j, anc); + return internal::read_as(_col_anc); + } + + iterator begin(int j, int anc, int chunk) const + { + return iterator(chunk, col_anc(j, anc)); + } + + iterator begin(int j, int anc) const + { + return begin(j, anc, 0); + } + + iterator end(int j, int anc) const + { + return begin(j, anc, n_chunks(j, anc)); + } + + rowarr_value_t to_dense( + size_t n_threads + ) const; + + std::tuple> write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads + ) const; + + // Overload: write only a selected subset of ancestries. + // - A_total is the total number of ancestries present in the input labels + // (upper bound; used for validation). + // - selected_ancestries lists the ancestry labels (from the 0..A_total-1 + // space) to be serialized. The output will have (1 + selected.size()) + // columns per SNP block. + std::tuple> write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const std::vector& selected_ancestries, + size_t n_threads + ) const; +}; + +} // namespace io +} // namespace adelie_core \ No newline at end of file diff --git a/adelie/src/include/adelie_core/io/io_snp_combine_r.ipp b/adelie/src/include/adelie_core/io/io_snp_combine_r.ipp new file mode 100644 index 00000000..c58bf205 --- /dev/null +++ b/adelie/src/include/adelie_core/io/io_snp_combine_r.ipp @@ -0,0 +1,814 @@ +#pragma once +#include +#include +#include +#include + +namespace adelie_core { +namespace io { + +ADELIE_CORE_IO_SNP_COMBINE_R_TP +size_t +ADELIE_CORE_IO_SNP_COMBINE_R::read() +{ + const size_t total_bytes = base_t::read(); + + size_t idx = sizeof(bool_t); + + _rows = internal::read_as(_buffer.data() + idx); + idx += sizeof(outer_t); + + _snps = internal::read_as(_buffer.data() + idx); + idx += sizeof(outer_t); + + _ancestries = internal::read_as(_buffer.data() + idx); + idx += sizeof(chunk_inner_t); + + _cols = _snps * (1 + _ancestries); + + _nnz.resize(_cols); + std::memcpy(_nnz.data(), _buffer.data() + idx, sizeof(outer_t) * _cols); + idx += sizeof(outer_t) * _cols; + + _outer.resize(_snps + 1); + std::memcpy(_outer.data(), _buffer.data() + idx, sizeof(outer_t) * (_snps + 1)); + idx += sizeof(outer_t) * (_snps + 1); + + return total_bytes; +} + +ADELIE_CORE_IO_SNP_COMBINE_R_TP +typename ADELIE_CORE_IO_SNP_COMBINE_R::rowarr_value_t +ADELIE_CORE_IO_SNP_COMBINE_R::to_dense( + size_t n_threads +) const +{ + const size_t n = rows(); + const size_t s = snps(); + const size_t A = ancestries(); + rowarr_value_t dense(n, s * (1 + A)); + + const auto routine = [&](outer_t j) { + auto dense_j = dense.middleCols((1 + A) * j, 1 + A); + dense_j.setZero(); + + // First column: SNP data + auto it = this->begin(j, 0); + const auto end = this->end(j, 0); + for (; it != end; ++it) { + dense_j(*it, 0) += 1; + } + + // Next A columns: Ancestry dosages + for (size_t a = 0; a < A; ++a) { + auto it = this->begin(j, a + 1); + const auto end = this->end(j, a + 1); + for (; it != end; ++it) { + dense_j(*it, a + 1) += 1; + } + } + }; + util::omp_parallel_for(routine, 0, s, n_threads); + + return dense; +} + +ADELIE_CORE_IO_SNP_COMBINE_R_TP +std::tuple> +ADELIE_CORE_IO_SNP_COMBINE_R::write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads +) const +{ + using sw_t = util::Stopwatch; + + if ( + (calldata.rows() != ancestries.rows()) || + (2 * calldata.cols() != ancestries.cols()) + ) { + throw util::adelie_core_error( + "ancestries must have shape (n, 2*s) where calldata has shape (n,s)." + ); + } + + if (A >= chunk_size) { + throw util::adelie_core_error( + "Number of ancestries A must be < " + + std::to_string(chunk_size) + + "." + ); + } + + sw_t sw; + std::unordered_map benchmark; + + const bool_t endian = is_big_endian(); + const outer_t n = calldata.rows(); + const outer_t s = calldata.cols(); + + const size_t max_chunks = (n + chunk_size - 1) / chunk_size; + if (max_chunks >= _max_inner) { + throw util::adelie_core_error( + "calldata dimensions are too large! " + ); + } + + // ------------------------------------------------------------------ + // 1. Compute sample_nnz (for file header) and payload_nnz separately + // ------------------------------------------------------------------ + vec_outer_t sample_nnz(s * (1 + A)); + vec_outer_t payload_nnz(s * (1 + A)); + sample_nnz.setZero(); + payload_nnz.setZero(); + sw.start(); + + // — SNP data column — + // sample_nnz = count of rows where calldata != 0 + // payload_nnz = sum of calldata values (0,1,2) + for (size_t j = 0; j < s; ++j) { + const auto cal_j = calldata.col(j); + outer_t samples = 0, bytes = 0; + for (outer_t i = 0; i < n; ++i) { + auto v = cal_j.coeff(i); + if (v != 0) ++samples; + bytes += static_cast(v); + } + size_t col_idx = (1 + A) * j; + sample_nnz[col_idx] = samples; + payload_nnz[col_idx] = bytes; + } + + // — Ancestry dosage columns — + // sample_nnz = count of rows where at least one haplo==a + // payload_nnz = count of *all* matching haplotypes (0–2) + for (size_t a = 0; a < A; ++a) { + for (size_t j = 0; j < s; ++j) { + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0, bytes = 0; + for (outer_t i = 0; i < n; ++i) { + bool hit0 = (anc0.coeff(i) == static_cast(a)); + bool hit1 = (anc1.coeff(i) == static_cast(a)); + if (hit0 || hit1) ++samples; + if (hit0) ++bytes; + if (hit1) ++bytes; + } + size_t col_idx = (1 + A) * j + (a + 1); + sample_nnz[col_idx] = samples; + payload_nnz[col_idx] = bytes; + } + } + benchmark["nnz"] = sw.elapsed(); + + + + // allocate sufficient memory (upper bound on size) + const size_t preamble_size = ( + sizeof(bool_t) + // endian + 2 * sizeof(outer_t) + // n, s + sizeof(chunk_inner_t) + // A + payload_nnz.size() * sizeof(outer_t) + // payload_nnz + (s + 1) * sizeof(outer_t) // outer (snps) + ); + buffer_t buffer( + preamble_size + + // for each SNP, reserve space for its (1 + A) columns structure: + s * ( + (1 + A) * sizeof(outer_t) + // outer pointers for each column + (1 + A) * ( + sizeof(inner_t) + // n_chunks + max_chunks * ( // for each chunk + sizeof(inner_t) + // chunk index + sizeof(chunk_inner_t) // chunk nnz - 1 + ) + ) + ) + + payload_nnz.sum() * sizeof(chunk_inner_t) // actual data entries + ); + + // populate buffer + outer_t idx = 0; + std::memcpy(buffer.data()+idx, &endian, sizeof(bool_t)); idx += sizeof(bool_t); + std::memcpy(buffer.data()+idx, &n, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &s, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &A, sizeof(chunk_inner_t)); idx += sizeof(chunk_inner_t); + std::memcpy(buffer.data()+idx, sample_nnz.data(), sizeof(outer_t) * sample_nnz.size()); + idx += sizeof(outer_t) * sample_nnz.size(); + + // outer[i] = number of bytes to jump from beginning of file + // to start reading snp i. + // outer[i+1] - outer[i] = total number of bytes for snp i. + char* const outer_ptr = buffer.data() + idx; + const size_t outer_size = s + 1; + idx += sizeof(outer_t) * outer_size; + std::memcpy(outer_ptr, &idx, sizeof(outer_t)); + + // flag to detect any errors + std::atomic_char try_failed = 0; + + // populate outer + const auto outer_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + + // Start with the space for (1+A) outer pointers + outer_t snp_bytes = (1 + A) * sizeof(outer_t); + + // --- SNP data column (col 0) --- + snp_bytes += sizeof(inner_t); // n_chunks + const auto cal_j = calldata.col(j); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + // count how many bytes we'll emit in this chunk + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + // each non‐zero allele produces one byte + nnz += cal_j[didx]; + } + // chunk header (index + nnz-1) if nonempty + if (nnz > 0) { + snp_bytes += sizeof(inner_t) // chunk index + + sizeof(chunk_inner_t) // chunk_nnz + + nnz * sizeof(chunk_inner_t); + } + } + + // --- Ancestry dosage columns (cols 1…A) --- + const auto anc_j0 = ancestries.col(2*j); + const auto anc_j1 = ancestries.col(2*j + 1); + for (size_t a = 0; a < A; ++a) { + snp_bytes += sizeof(inner_t); // n_chunks for this ancestry + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + // one entry per haplotype + if (anc_j0[didx] == static_cast(a)) ++nnz; + if (anc_j1[didx] == static_cast(a)) ++nnz; + } + if (nnz > 0) { + snp_bytes += sizeof(inner_t) // chunk index + + sizeof(chunk_inner_t) // chunk_nnz + + nnz * sizeof(chunk_inner_t); + } + } + } + + // write the per-SNP byte count into the outer table + std::memcpy( + outer_ptr + sizeof(outer_t) * (j+1), + &snp_bytes, + sizeof(outer_t) + ); + }; + sw.start(); + util::omp_parallel_for(outer_routine, 0, s, n_threads); + benchmark["outer_time"] = sw.elapsed(); + + switch (try_failed) { + case 1: { + throw util::adelie_core_error( + "Detected an ancestry not in the range [0, A). " + "Make sure ancestries only contains values in [0, A). " + ); + break; + } + case 2: { + throw util::adelie_core_error( + "Detected a value not in [0,2]. " + "Make sure calldata only contains values in [0,2]. " + ); + break; + } + } + + // cumsum outer + for (outer_t j = 0; j < s; ++j) { + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t) * (j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t) * j); + const outer_t sum = outer_curr + outer_prev; + std::memcpy(outer_ptr + sizeof(outer_t) * (j+1), &sum, sizeof(outer_t)); + } + + const outer_t outer_last = internal::read_as(outer_ptr + sizeof(outer_t) * s); + + if (outer_last > static_cast(buffer.size())) { + throw util::adelie_core_error( + "Buffer was not initialized with a large enough size. " + "\n\tBuffer size: " + std::to_string(buffer.size()) + + "\n\tExpected size: " + std::to_string(outer_last) + + "\nThis is likely a bug in the code. Please report it! " + ); + } + idx = outer_last; + + // populate (column) inner buffers + try_failed = 0; + const auto inner_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + + // 1) Slice out this SNP's block + const outer_t outer_curr = internal::read_as( + outer_ptr + sizeof(outer_t)*(j+1) + ); + const outer_t outer_prev = internal::read_as( + outer_ptr + sizeof(outer_t)*j + ); + Eigen::Map buffer_j( + buffer.data() + outer_prev, + outer_curr - outer_prev + ); + + // 2) Skip over the (1+A) outer pointers + outer_t cidx = (1 + A) * sizeof(outer_t); + + // + // --- SNP data column (col 0 of the block) --- + // + // write pointer to our chunk data + std::memcpy(buffer_j.data() + 0*sizeof(outer_t), &cidx, sizeof(outer_t)); + // reserve space for 'n_chunks' + auto *n_chunks_ptr0 = reinterpret_cast(buffer_j.data() + cidx); + cidx += sizeof(inner_t); + inner_t n_chunks0 = 0; + + // grab the whole column once + const auto cal_j = calldata.col(j); + + // for each chunk + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + + // chunk header + auto *chunk_index = reinterpret_cast(buffer_j.data() + curr_idx); + curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast( + buffer_j.data() + curr_idx + ); + curr_idx += sizeof(chunk_inner_t); + + // now list out positions, *once per allele* + auto *chunk_begin = reinterpret_cast( + buffer_j.data() + curr_idx + ); + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + // if cal_j[didx] is 1 or 2, emit that many entries + for (int rep = 0; rep < cal_j[didx]; ++rep) { + chunk_begin[nnz++] = c; + curr_idx += sizeof(chunk_inner_t); + } + } + + if (nnz) { + // commit header and advance + std::memcpy(chunk_index, &k, sizeof(inner_t)); + *chunk_nnz = nnz - 1; + cidx = curr_idx; + ++n_chunks0; + } + } + // write back total chunk‐count + std::memcpy(n_chunks_ptr0, &n_chunks0, sizeof(inner_t)); + + + // + // --- Ancestry dosage columns (cols 1…A of the block) --- + // + // grab both haplotype‐columns once + const auto anc_j0 = ancestries.col(2*j); + const auto anc_j1 = ancestries.col(2*j + 1); + + for (size_t a = 0; a < A; ++a) { + // pointer for this ancestry col + std::memcpy( + buffer_j.data() + (1 + a)*sizeof(outer_t), + &cidx, + sizeof(outer_t) + ); + auto *n_chunks_ptr = reinterpret_cast( + buffer_j.data() + cidx + ); + cidx += sizeof(inner_t); + inner_t n_chunks = 0; + + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + + auto *chunk_index = reinterpret_cast( + buffer_j.data() + curr_idx + ); + curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast( + buffer_j.data() + curr_idx + ); + curr_idx += sizeof(chunk_inner_t); + + auto *chunk_begin = reinterpret_cast( + buffer_j.data() + curr_idx + ); + inner_t nnz = 0; + + // emit one entry for each haplotype of ancestry 'a' + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if (anc_j0[didx] == static_cast(a)) { + chunk_begin[nnz++] = c; + curr_idx += sizeof(chunk_inner_t); + } + if (anc_j1[didx] == static_cast(a)) { + chunk_begin[nnz++] = c; + curr_idx += sizeof(chunk_inner_t); + } + } + + if (nnz) { + std::memcpy(chunk_index, &k, sizeof(inner_t)); + *chunk_nnz = nnz - 1; + cidx = curr_idx; + ++n_chunks; + } + } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + + // final size check + if (cidx != static_cast(buffer_j.size())) { + try_failed = 1; + } + }; + + sw.start(); + util::omp_parallel_for(inner_routine, 0, s, n_threads); + benchmark["inner"] = sw.elapsed(); + + if (try_failed) { + throw util::adelie_core_error( + "Column index certificate does not match expected size. " + "This is likely a bug in the code. Please report it! " + ); + } + + sw.start(); + auto file_ptr = fopen_safe(_filename.c_str(), "wb"); + auto fp = file_ptr.get(); + auto total_bytes = std::fwrite(buffer.data(), sizeof(char), idx, fp); + if (total_bytes != static_cast(idx)) { + throw util::adelie_core_error( + "Could not write the full buffer." + ); + } + benchmark["write"] = sw.elapsed(); + + return {total_bytes, benchmark}; +} + +} // namespace io +} // namespace adelie_core + +// ---------------------------------------------------------------------------- +// Overload: write with selected ancestry list +// ---------------------------------------------------------------------------- +#include + +namespace adelie_core { +namespace io { + +ADELIE_CORE_IO_SNP_COMBINE_R_TP +std::tuple> +ADELIE_CORE_IO_SNP_COMBINE_R::write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const std::vector& selected_ancestries, + size_t n_threads +) const +{ + using sw_t = util::Stopwatch; + + if ( + (calldata.rows() != ancestries.rows()) || + (2 * calldata.cols() != ancestries.cols()) + ) { + throw util::adelie_core_error( + "ancestries must have shape (n, 2*s) where calldata has shape (n,s)." + ); + } + + // Validate selection + if (selected_ancestries.empty()) { + throw util::adelie_core_error( + "selected_ancestries must be non-empty." + ); + } + for (auto a : selected_ancestries) { + if (a >= A_total) { + throw util::adelie_core_error( + "selected_ancestries contains an index out of range [0, A_total)." + ); + } + } + + // Ensure entries are unique while preserving the user-provided order + std::vector sel; + sel.reserve(selected_ancestries.size()); + for (auto a : selected_ancestries) { + if (std::find(sel.begin(), sel.end(), a) == sel.end()) { + sel.push_back(a); + } + } + + // Now proceed similarly to the A-based writer, with A := sel.size() and + // ancestry references restricted to 'sel'. Keep A_total only for validation. + + const bool_t endian = is_big_endian(); + const outer_t n = calldata.rows(); + const outer_t s = calldata.cols(); + const size_t A_sel = sel.size(); + + if (A_sel >= chunk_size) { + throw util::adelie_core_error( + "Number of selected ancestries must be < " + + std::to_string(chunk_size) + + "." + ); + } + + const size_t max_chunks = (n + chunk_size - 1) / chunk_size; + if (max_chunks >= _max_inner) { + throw util::adelie_core_error( + "calldata dimensions are too large! " + ); + } + + sw_t sw; + std::unordered_map benchmark; + + // nnz counts for header (samples with nonzero per column of output) + vec_outer_t sample_nnz(s * (1 + A_sel)); + vec_outer_t payload_nnz(s * (1 + A_sel)); + sample_nnz.setZero(); + payload_nnz.setZero(); + + sw.start(); + // SNP data column stats + for (size_t j = 0; j < s; ++j) { + const auto cal_j = calldata.col(j); + outer_t samples = 0, bytes = 0; + for (outer_t i = 0; i < n; ++i) { + auto v = cal_j.coeff(i); + if (v != 0) ++samples; + bytes += static_cast(v); + } + size_t col_idx = (1 + A_sel) * j; + sample_nnz[col_idx] = samples; + payload_nnz[col_idx] = bytes; + } + + // Ancestry dosage columns stats (restricted to sel) + for (size_t ai = 0; ai < A_sel; ++ai) { + const auto a = sel[ai]; + for (size_t j = 0; j < s; ++j) { + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0, bytes = 0; + for (outer_t i = 0; i < n; ++i) { + bool hit0 = (anc0.coeff(i) == static_cast(a)); + bool hit1 = (anc1.coeff(i) == static_cast(a)); + if (hit0 || hit1) ++samples; + if (hit0) ++bytes; + if (hit1) ++bytes; + } + size_t col_idx = (1 + A_sel) * j + (ai + 1); + sample_nnz[col_idx] = samples; + payload_nnz[col_idx] = bytes; + } + } + benchmark["nnz"] = sw.elapsed(); + + // Buffer pre-allocation + const size_t preamble_size = ( + sizeof(bool_t) + // endian + 2 * sizeof(outer_t) + // n, s + sizeof(chunk_inner_t) + // A_sel (encoded as chunk_inner_t) + sample_nnz.size() * sizeof(outer_t) + // nnz per output column + (s + 1) * sizeof(outer_t) // outer (snps) + ); + buffer_t buffer( + preamble_size + + s * ( + (1 + A_sel) * sizeof(outer_t) + + (1 + A_sel) * ( + sizeof(inner_t) + + max_chunks * ( + sizeof(inner_t) + + sizeof(chunk_inner_t) + ) + ) + ) + + payload_nnz.sum() * sizeof(chunk_inner_t) + ); + + // Populate header + outer_t idx = 0; + std::memcpy(buffer.data()+idx, &endian, sizeof(bool_t)); idx += sizeof(bool_t); + std::memcpy(buffer.data()+idx, &n, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &s, sizeof(outer_t)); idx += sizeof(outer_t); + const auto A_sel_byte = static_cast(A_sel); + std::memcpy(buffer.data()+idx, &A_sel_byte, sizeof(chunk_inner_t)); idx += sizeof(chunk_inner_t); + std::memcpy(buffer.data()+idx, sample_nnz.data(), sizeof(outer_t) * sample_nnz.size()); + idx += sizeof(outer_t) * sample_nnz.size(); + + // outer table + char* const outer_ptr = buffer.data() + idx; + const size_t outer_size = s + 1; + idx += sizeof(outer_t) * outer_size; + std::memcpy(outer_ptr, &idx, sizeof(outer_t)); + + std::atomic_char try_failed = 0; + + // Compute per-SNP byte sizes + const auto outer_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + + outer_t snp_bytes = (1 + A_sel) * sizeof(outer_t); + + // Genotype column + snp_bytes += sizeof(inner_t); + const auto cal_j = calldata.col(j); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + nnz += cal_j[didx]; + } + if (nnz > 0) { + snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz * sizeof(chunk_inner_t); + } + } + + // Selected ancestry columns + const auto anc_j0 = ancestries.col(2*j); + const auto anc_j1 = ancestries.col(2*j + 1); + for (size_t ai = 0; ai < A_sel; ++ai) { + const auto a = sel[ai]; + snp_bytes += sizeof(inner_t); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if (anc_j0[didx] == static_cast(a)) ++nnz; + if (anc_j1[didx] == static_cast(a)) ++nnz; + } + if (nnz > 0) { + snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz * sizeof(chunk_inner_t); + } + } + } + + std::memcpy( + outer_ptr + sizeof(outer_t) * (j+1), + &snp_bytes, + sizeof(outer_t) + ); + }; + sw.start(); + util::omp_parallel_for(outer_routine, 0, s, n_threads); + benchmark["outer_time"] = sw.elapsed(); + + // cumsum outer + for (outer_t j = 0; j < s; ++j) { + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t) * (j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t) * j); + const outer_t sum = outer_curr + outer_prev; + std::memcpy(outer_ptr + sizeof(outer_t) * (j+1), &sum, sizeof(outer_t)); + } + + const outer_t outer_last = internal::read_as(outer_ptr + sizeof(outer_t) * s); + if (outer_last > static_cast(buffer.size())) { + throw util::adelie_core_error( + "Buffer was not initialized with a large enough size. " + ); + } + idx = outer_last; + + // Fill per-SNP blocks + const auto inner_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t) * (j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t) * j); + Eigen::Map buffer_j( + buffer.data() + outer_prev, + outer_curr - outer_prev + ); + + outer_t cidx = (1 + A_sel) * sizeof(outer_t); + + // Genotype column pointer + std::memcpy(buffer_j.data() + 0*sizeof(outer_t), &cidx, sizeof(outer_t)); + auto *n_chunks_ptr0 = reinterpret_cast(buffer_j.data() + cidx); + cidx += sizeof(inner_t); + inner_t n_chunks0 = 0; + + const auto cal_j = calldata.col(j); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + auto *chunk_index = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(chunk_inner_t); + auto *chunk_begin = reinterpret_cast(buffer_j.data() + curr_idx); + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + for (int rep = 0; rep < cal_j[didx]; ++rep) { + chunk_begin[nnz++] = c; + curr_idx += sizeof(chunk_inner_t); + } + } + if (nnz) { + std::memcpy(chunk_index, &k, sizeof(inner_t)); + *chunk_nnz = nnz - 1; + cidx = curr_idx; + ++n_chunks0; + } + } + std::memcpy(n_chunks_ptr0, &n_chunks0, sizeof(inner_t)); + + // Selected ancestry columns + const auto anc_j0 = ancestries.col(2*j); + const auto anc_j1 = ancestries.col(2*j + 1); + for (size_t ai = 0; ai < A_sel; ++ai) { + const auto a = sel[ai]; + std::memcpy(buffer_j.data() + (1 + ai)*sizeof(outer_t), &cidx, sizeof(outer_t)); + auto *n_chunks_ptr = reinterpret_cast(buffer_j.data() + cidx); + cidx += sizeof(inner_t); + inner_t n_chunks = 0; + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + auto *chunk_index = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(chunk_inner_t); + auto *chunk_begin = reinterpret_cast(buffer_j.data() + curr_idx); + inner_t nnz = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if (anc_j0[didx] == static_cast(a)) { chunk_begin[nnz++] = c; curr_idx += sizeof(chunk_inner_t); } + if (anc_j1[didx] == static_cast(a)) { chunk_begin[nnz++] = c; curr_idx += sizeof(chunk_inner_t); } + } + if (nnz) { + std::memcpy(chunk_index, &k, sizeof(inner_t)); + *chunk_nnz = nnz - 1; + cidx = curr_idx; + ++n_chunks; + } + } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + + if (cidx != static_cast(buffer_j.size())) { + try_failed = 1; + } + }; + sw.start(); + util::omp_parallel_for(inner_routine, 0, s, n_threads); + benchmark["inner"] = sw.elapsed(); + + if (try_failed) { + throw util::adelie_core_error( + "Column index certificate does not match expected size. " + ); + } + + sw.start(); + auto file_ptr = fopen_safe(_filename.c_str(), "wb"); + auto fp = file_ptr.get(); + auto total_bytes = std::fwrite(buffer.data(), sizeof(char), idx, fp); + if (total_bytes != static_cast(idx)) { + throw util::adelie_core_error( + "Could not write the full buffer." + ); + } + benchmark["write"] = sw.elapsed(); + + return {total_bytes, benchmark}; +} + +} // namespace io +} // namespace adelie_core \ No newline at end of file diff --git a/adelie/src/include/adelie_core/io/io_snp_combine_s.hpp b/adelie/src/include/adelie_core/io/io_snp_combine_s.hpp new file mode 100644 index 00000000..3a0a8829 --- /dev/null +++ b/adelie/src/include/adelie_core/io/io_snp_combine_s.hpp @@ -0,0 +1,150 @@ +#pragma once +#include +#include + +#ifndef ADELIE_CORE_IO_SNP_COMBINE_S_TP +#define ADELIE_CORE_IO_SNP_COMBINE_S_TP \ + template +#endif +#ifndef ADELIE_CORE_IO_SNP_COMBINE_S +#define ADELIE_CORE_IO_SNP_COMBINE_S \ + IOSNPCombineS +#endif + +namespace adelie_core { +namespace io { + +template >> +class IOSNPCombineS : public IOSNPBase +{ +public: + using base_t = IOSNPBase; + using outer_t = uint64_t; + using inner_t = uint32_t; + using chunk_inner_t = uint8_t; + using value_t = int8_t; + using vec_outer_t = util::rowvec_type; + using vec_inner_t = util::rowvec_type; + using vec_value_t = util::rowvec_type; + using rowarr_value_t = util::rowarr_type; + using colarr_value_t = util::colarr_type; + using typename base_t::bool_t; + using typename base_t::buffer_t; + + static constexpr size_t n_bits_per_byte = 8; + // Use 128 chunk size to guarantee chunk_nnz fits in uint8 (nnz-1) + static constexpr size_t chunk_size = ( + static_cast(1UL) << (n_bits_per_byte * sizeof(chunk_inner_t)) + ) / 2; // 256/2 = 128 + +protected: + static constexpr size_t _max_inner = ( + static_cast(1UL) << (n_bits_per_byte * sizeof(inner_t)) + ); + + using base_t::throw_no_read; + using base_t::fopen_safe; + using base_t::is_big_endian; + using base_t::_buffer; + using base_t::_filename; + using base_t::_is_read; + + outer_t _rows; + outer_t _snps; + outer_t _ancestries; // number of ancestry labels A + outer_t _cols; // s * (2*A) + vec_outer_t _nnz; // per column sample nnz + vec_outer_t _outer; // per snp block offset + +public: + using iterator = IOSNPChunkIterator< + chunk_size, inner_t, chunk_inner_t + >; + + using base_t::base_t; + + size_t read() override; + + outer_t rows() const { + if (!_is_read) throw_no_read(); + return _rows; + } + + outer_t snps() const { + if (!_is_read) throw_no_read(); + return _snps; + } + + outer_t ancestries() const { + if (!_is_read) throw_no_read(); + return _ancestries; + } + + outer_t cols() const { + if (!_is_read) throw_no_read(); + return _cols; + } + + Eigen::Ref nnz() const { + if (!_is_read) throw_no_read(); + return _nnz; + } + + Eigen::Ref outer() const { + if (!_is_read) throw_no_read(); + return _outer; + } + + Eigen::Ref col(int j) const { + return Eigen::Map( + _buffer.data() + _outer[j], + _outer[j+1] - _outer[j] + ); + } + + // Column within block: 0..A-1 => mutated-haplotypes-per-ancestry + // A..2A-1 => ancestry-dosage-per-ancestry + const char* col_k(int j, int k) const { + const auto _col = col(j); + return ( + _col.data() + + internal::read_as(_col.data() + sizeof(outer_t) * k) + ); + } + + inner_t n_chunks(int j, int k) const { + const auto* _col_k = col_k(j, k); + return internal::read_as(_col_k); + } + + iterator begin(int j, int k, int chunk) const { + return iterator(chunk, col_k(j, k)); + } + + iterator begin(int j, int k) const { return begin(j, k, 0); } + iterator end(int j, int k) const { return begin(j, k, n_chunks(j, k)); } + + rowarr_value_t to_dense( + size_t n_threads + ) const; + + std::tuple> write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads + ) const; + + std::tuple> write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const std::vector& selected_ancestries, + size_t n_threads + ) const; +}; + +} // namespace io +} // namespace adelie_core + + diff --git a/adelie/src/include/adelie_core/io/io_snp_combine_s.ipp b/adelie/src/include/adelie_core/io/io_snp_combine_s.ipp new file mode 100644 index 00000000..ae251e7c --- /dev/null +++ b/adelie/src/include/adelie_core/io/io_snp_combine_s.ipp @@ -0,0 +1,559 @@ +#pragma once +#include +#include +#include +#include +#include + +namespace adelie_core { +namespace io { + +ADELIE_CORE_IO_SNP_COMBINE_S_TP +size_t +ADELIE_CORE_IO_SNP_COMBINE_S::read() +{ + const size_t total_bytes = base_t::read(); + + size_t idx = sizeof(bool_t); + + _rows = internal::read_as(_buffer.data() + idx); + idx += sizeof(outer_t); + + _snps = internal::read_as(_buffer.data() + idx); + idx += sizeof(outer_t); + + _ancestries = internal::read_as(_buffer.data() + idx); + idx += sizeof(chunk_inner_t); + + _cols = _snps * (2 * _ancestries); + + _nnz.resize(_cols); + std::memcpy(_nnz.data(), _buffer.data() + idx, sizeof(outer_t) * _cols); + idx += sizeof(outer_t) * _cols; + + _outer.resize(_snps + 1); + std::memcpy(_outer.data(), _buffer.data() + idx, sizeof(outer_t) * (_snps + 1)); + idx += sizeof(outer_t) * (_snps + 1); + + return total_bytes; +} + +ADELIE_CORE_IO_SNP_COMBINE_S_TP +typename ADELIE_CORE_IO_SNP_COMBINE_S::rowarr_value_t +ADELIE_CORE_IO_SNP_COMBINE_S::to_dense( + size_t n_threads +) const +{ + const size_t n = rows(); + const size_t s = snps(); + const size_t A = ancestries(); + rowarr_value_t dense(n, s * (2 * A)); + + const auto routine = [&](outer_t j) { + auto dense_j = dense.middleCols((2 * A) * j, 2 * A); + dense_j.setZero(); + + // First A columns: mutated haplotype counts per ancestry + for (size_t a = 0; a < A; ++a) { + auto it = this->begin(j, static_cast(a)); + const auto end = this->end(j, static_cast(a)); + for (; it != end; ++it) { + dense_j(*it, a) += 1; + } + } + + // Next A columns: ancestry dosage counts per ancestry (like unphased ancestry dosage) + for (size_t a = 0; a < A; ++a) { + auto it = this->begin(j, static_cast(A + a)); + const auto end = this->end(j, static_cast(A + a)); + for (; it != end; ++it) { + dense_j(*it, A + a) += 1; + } + } + }; + util::omp_parallel_for(routine, 0, s, n_threads); + + return dense; +} + +ADELIE_CORE_IO_SNP_COMBINE_S_TP +std::tuple> +ADELIE_CORE_IO_SNP_COMBINE_S::write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads +) const +{ + using sw_t = util::Stopwatch; + + // calldata and ancestries are phased: shape (n, 2*s) + if ( + (calldata.rows() != ancestries.rows()) || + (calldata.cols() != ancestries.cols()) || + (calldata.cols() % 2) + ) { + throw util::adelie_core_error( + "calldata and ancestries must have shape (n, 2*s)." + ); + } + + if (A >= chunk_size) { + throw util::adelie_core_error( + "Number of ancestries A must be < " + + std::to_string(chunk_size) + + "." + ); + } + + sw_t sw; + std::unordered_map benchmark; + + const bool_t endian = is_big_endian(); + const outer_t n = calldata.rows(); + const outer_t s = calldata.cols() / 2; + + const size_t max_chunks = (n + chunk_size - 1) / chunk_size; + if (max_chunks >= _max_inner) { + throw util::adelie_core_error( + "calldata dimensions are too large! " + ); + } + + // nnz per output column (2*A per SNP) + vec_outer_t nnz(s * (2 * A)); + nnz.setZero(); + sw.start(); + // First A: mutated hap counts (sum over haps) + for (size_t a = 0; a < A; ++a) { + for (size_t j = 0; j < s; ++j) { + const auto cal0 = calldata.col(2*j); + const auto cal1 = calldata.col(2*j+1); + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0; + for (outer_t i = 0; i < n; ++i) { + bool contribute = ( + (cal0[i] == 1 && anc0[i] == static_cast(a)) || + (cal1[i] == 1 && anc1[i] == static_cast(a)) + ); + if (contribute) ++samples; // sample nnz counts rows with any hit + } + nnz[(2*A)*j + a] = samples; + } + } + // Next A: ancestry dosage counts + for (size_t a = 0; a < A; ++a) { + for (size_t j = 0; j < s; ++j) { + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0; + for (outer_t i = 0; i < n; ++i) { + if ((anc0[i] == static_cast(a)) || (anc1[i] == static_cast(a))) ++samples; + } + nnz[(2*A)*j + (A + a)] = samples; + } + } + benchmark["nnz"] = sw.elapsed(); + + // Pre-allocate buffer + const size_t preamble_size = ( + sizeof(bool_t) + + 2 * sizeof(outer_t) + + sizeof(chunk_inner_t) + + nnz.size() * sizeof(outer_t) + + (s + 1) * sizeof(outer_t) + ); + buffer_t buffer( + preamble_size + + s * ( + (2*A) * sizeof(outer_t) + + (2*A) * ( + sizeof(inner_t) + + max_chunks * ( + sizeof(inner_t) + + sizeof(chunk_inner_t) + ) + ) + ) + + // payload: number of actual entries written per column equals + // mutated hap counts (each hit writes 1) and dosage counts (each hap hit writes 1) + // Use nnz as an upper bound for sample nnz; worst-case payload could be larger for dosage. + // To keep it simple and safe, over-allocate by factor 2. + 2 * nnz.sum() * sizeof(chunk_inner_t) + ); + + // Header + outer_t idx = 0; + std::memcpy(buffer.data()+idx, &endian, sizeof(bool_t)); idx += sizeof(bool_t); + std::memcpy(buffer.data()+idx, &n, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &s, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &A, sizeof(chunk_inner_t)); idx += sizeof(chunk_inner_t); + std::memcpy(buffer.data()+idx, nnz.data(), sizeof(outer_t) * nnz.size()); + idx += sizeof(outer_t) * nnz.size(); + + // outer table (per SNP block) + char* const outer_ptr = buffer.data() + idx; + idx += sizeof(outer_t) * (s + 1); + std::memcpy(outer_ptr, &idx, sizeof(outer_t)); + + std::atomic_char try_failed = 0; + + // Outer routine: compute per-SNP byte sizes + const auto outer_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + outer_t snp_bytes = (2*A) * sizeof(outer_t); + + const auto cal0 = calldata.col(2*j); + const auto cal1 = calldata.col(2*j+1); + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + + // First A mutated-hap columns (write one entry per mutated hap) + for (size_t a = 0; a < A; ++a) { + snp_bytes += sizeof(inner_t); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + // Validate + if ((cal0[didx] != 0 && cal0[didx] != 1) || (cal1[didx] != 0 && cal1[didx] != 1)) { try_failed = 2; return; } + if (anc0[didx] < 0 || anc0[didx] >= static_cast(A) || anc1[didx] < 0 || anc1[didx] >= static_cast(A)) { try_failed = 1; return; } + if (cal0[didx] == 1 && anc0[didx] == static_cast(a)) ++nnz_chunk; + if (cal1[didx] == 1 && anc1[didx] == static_cast(a)) ++nnz_chunk; + } + if (nnz_chunk > 0) { + snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz_chunk * sizeof(chunk_inner_t); + } + } + } + + // Next A dosage columns + for (size_t a = 0; a < A; ++a) { + snp_bytes += sizeof(inner_t); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if (anc0[didx] < 0 || anc0[didx] >= static_cast(A) || anc1[didx] < 0 || anc1[didx] >= static_cast(A)) { try_failed = 1; return; } + if (anc0[didx] == static_cast(a)) ++nnz_chunk; + if (anc1[didx] == static_cast(a)) ++nnz_chunk; + } + if (nnz_chunk > 0) { + snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz_chunk * sizeof(chunk_inner_t); + } + } + } + + std::memcpy(outer_ptr + sizeof(outer_t) * (j+1), &snp_bytes, sizeof(outer_t)); + }; + util::omp_parallel_for(outer_routine, 0, s, n_threads); + switch (try_failed) { + case 1: { + throw util::adelie_core_error( + "Detected an ancestry not in the range [0, A). Make sure ancestries only contains values in [0, A). " + ); + } + case 2: { + throw util::adelie_core_error( + "Detected a non-binary value. Make sure calldata only contains 0 or 1 values. " + ); + } + } + + // cumsum outer + for (outer_t j = 0; j < s; ++j) { + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t) * (j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t) * j); + const outer_t sum = outer_curr + outer_prev; + std::memcpy(outer_ptr + sizeof(outer_t) * (j+1), &sum, sizeof(outer_t)); + } + const outer_t outer_last = internal::read_as(outer_ptr + sizeof(outer_t) * s); + if (outer_last > static_cast(buffer.size())) { + throw util::adelie_core_error("Buffer was not initialized with a large enough size. "); + } + idx = outer_last; + + // Inner routine: write per-SNP block detail + const auto inner_routine = [&](outer_t j) { + if (try_failed.load(std::memory_order_relaxed)) return; + + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t) * (j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t) * j); + Eigen::Map buffer_j( + buffer.data() + outer_prev, + outer_curr - outer_prev + ); + + outer_t cidx = (2*A) * sizeof(outer_t); + + const auto cal0 = calldata.col(2*j); + const auto cal1 = calldata.col(2*j+1); + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + + // First A: mutated hap columns (emit one entry per mutated hap) + for (size_t a = 0; a < A; ++a) { + std::memcpy(buffer_j.data() + sizeof(outer_t) * a, &cidx, sizeof(outer_t)); + auto *n_chunks_ptr = reinterpret_cast(buffer_j.data() + cidx); + cidx += sizeof(inner_t); + inner_t n_chunks = 0; + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + auto *chunk_index = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(chunk_inner_t); + auto *chunk_begin = reinterpret_cast(buffer_j.data() + curr_idx); + inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if ((cal0[didx] != 0 && cal0[didx] != 1) || (cal1[didx] != 0 && cal1[didx] != 1)) { try_failed = 2; return; } + if (anc0[didx] < 0 || anc0[didx] >= static_cast(A) || anc1[didx] < 0 || anc1[didx] >= static_cast(A)) { try_failed = 1; return; } + if (cal0[didx] == 1 && anc0[didx] == static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + if (cal1[didx] == 1 && anc1[didx] == static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + } + if (nnz_chunk) { std::memcpy(chunk_index, &k, sizeof(inner_t)); *chunk_nnz = nnz_chunk - 1; cidx = curr_idx; ++n_chunks; } + } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + + // Next A: ancestry dosage columns + for (size_t a = 0; a < A; ++a) { + std::memcpy(buffer_j.data() + sizeof(outer_t) * (A + a), &cidx, sizeof(outer_t)); + auto *n_chunks_ptr = reinterpret_cast(buffer_j.data() + cidx); + cidx += sizeof(inner_t); + inner_t n_chunks = 0; + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; + size_t curr_idx = cidx; + auto *chunk_index = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast(buffer_j.data() + curr_idx); curr_idx += sizeof(chunk_inner_t); + auto *chunk_begin = reinterpret_cast(buffer_j.data() + curr_idx); + inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; + if (didx >= n) break; + if (anc0[didx] == static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + if (anc1[didx] == static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + } + if (nnz_chunk) { std::memcpy(chunk_index, &k, sizeof(inner_t)); *chunk_nnz = nnz_chunk - 1; cidx = curr_idx; ++n_chunks; } + } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + + if (cidx != static_cast(buffer_j.size())) { try_failed = 3; } + }; + util::omp_parallel_for(inner_routine, 0, s, n_threads); + + switch (try_failed) { + case 1: throw util::adelie_core_error("Detected an ancestry not in the range [0, A). "); + case 2: throw util::adelie_core_error("Detected a non-binary value. "); + case 3: throw util::adelie_core_error("Column index certificate does not match expected size. "); + } + + auto file_ptr = fopen_safe(_filename.c_str(), "wb"); + auto fp = file_ptr.get(); + auto total_bytes = std::fwrite(buffer.data(), sizeof(char), idx, fp); + if (total_bytes != static_cast(idx)) { + throw util::adelie_core_error("Could not write the full buffer."); + } + benchmark["write"] = 0.0; // minimal tracking here + + return {total_bytes, benchmark}; +} + +ADELIE_CORE_IO_SNP_COMBINE_S_TP +std::tuple> +ADELIE_CORE_IO_SNP_COMBINE_S::write( + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const std::vector& selected_ancestries, + size_t n_threads +) const +{ + // Implement by filtering ancestries and delegating to the non-selected write + if (selected_ancestries.empty()) { + throw util::adelie_core_error("selected_ancestries must be non-empty."); + } + for (auto a : selected_ancestries) { + if (a >= A_total) { + throw util::adelie_core_error("selected_ancestries contains an index out of range [0, A_total)."); + } + } + // Ensure uniqueness preserving order + std::vector sel; + sel.reserve(selected_ancestries.size()); + for (auto a : selected_ancestries) { + if (std::find(sel.begin(), sel.end(), a) == sel.end()) sel.push_back(a); + } + + const outer_t n = calldata.rows(); + const outer_t s = calldata.cols() / 2; + const size_t B = sel.size(); + + // Build temporary dense then serialize via the dense path: reuse the same writer logic by constructing columns directly + // For efficiency, we implement selected-version by emitting directly, mirroring the unselected write with A := B and ancestry index remapped via sel. + + using sw_t = util::Stopwatch; + sw_t sw; + std::unordered_map benchmark; + + const bool_t endian = is_big_endian(); + const size_t max_chunks = (n + chunk_size - 1) / chunk_size; + if (max_chunks >= _max_inner) { + throw util::adelie_core_error("calldata dimensions are too large! "); + } + + vec_outer_t nnz(s * (2 * B)); nnz.setZero(); + // nnz for mutated-hap columns + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + for (size_t j = 0; j < s; ++j) { + const auto cal0 = calldata.col(2*j); + const auto cal1 = calldata.col(2*j+1); + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0; + for (outer_t i = 0; i < n; ++i) { + bool contribute = ( + (cal0[i] == 1 && anc0[i] == static_cast(a)) || + (cal1[i] == 1 && anc1[i] == static_cast(a)) + ); + if (contribute) ++samples; + } + nnz[(2*B)*j + bi] = samples; + } + } + // nnz for dosage columns + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + for (size_t j = 0; j < s; ++j) { + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + outer_t samples = 0; + for (outer_t i = 0; i < n; ++i) { + if ((anc0[i] == static_cast(a)) || (anc1[i] == static_cast(a))) ++samples; + } + nnz[(2*B)*j + (B + bi)] = samples; + } + } + + const outer_t n_cols = s * (2 * B); + const size_t preamble_size = ( + sizeof(bool_t) + 2*sizeof(outer_t) + sizeof(chunk_inner_t) + n_cols*sizeof(outer_t) + (s+1)*sizeof(outer_t) + ); + buffer_t buffer( + preamble_size + + s * ( + (2*B) * sizeof(outer_t) + + (2*B) * ( sizeof(inner_t) + max_chunks * (sizeof(inner_t) + sizeof(chunk_inner_t)) ) + ) + 2 * nnz.sum() * sizeof(chunk_inner_t) + ); + + outer_t idx = 0; + std::memcpy(buffer.data()+idx, &endian, sizeof(bool_t)); idx += sizeof(bool_t); + std::memcpy(buffer.data()+idx, &n, sizeof(outer_t)); idx += sizeof(outer_t); + std::memcpy(buffer.data()+idx, &s, sizeof(outer_t)); idx += sizeof(outer_t); + const auto B_byte = static_cast(B); + std::memcpy(buffer.data()+idx, &B_byte, sizeof(chunk_inner_t)); idx += sizeof(chunk_inner_t); + std::memcpy(buffer.data()+idx, nnz.data(), sizeof(outer_t) * nnz.size()); idx += sizeof(outer_t) * nnz.size(); + + char* const outer_ptr = buffer.data() + idx; idx += sizeof(outer_t) * (s+1); std::memcpy(outer_ptr, &idx, sizeof(outer_t)); + + const auto outer_routine = [&](outer_t j) { + outer_t snp_bytes = (2*B) * sizeof(outer_t); + const auto cal0 = calldata.col(2*j); + const auto cal1 = calldata.col(2*j+1); + const auto anc0 = ancestries.col(2*j); + const auto anc1 = ancestries.col(2*j+1); + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + snp_bytes += sizeof(inner_t); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; if (didx >= n) break; + bool h0 = (cal0[didx]==1 && anc0[didx]==static_cast(a)); + bool h1 = (cal1[didx]==1 && anc1[didx]==static_cast(a)); + if (h0) ++nnz_chunk; + if (h1) ++nnz_chunk; + } + if (nnz_chunk) snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz_chunk * sizeof(chunk_inner_t); + } + } + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + snp_bytes += sizeof(inner_t); + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { const outer_t didx = base + c; if (didx >= n) break; if (anc0[didx]==static_cast(a)) ++nnz_chunk; if (anc1[didx]==static_cast(a)) ++nnz_chunk; } + if (nnz_chunk) snp_bytes += sizeof(inner_t) + sizeof(chunk_inner_t) + nnz_chunk * sizeof(chunk_inner_t); + } + } + std::memcpy(outer_ptr + sizeof(outer_t)*(j+1), &snp_bytes, sizeof(outer_t)); + }; + util::omp_parallel_for(outer_routine, 0, s, n_threads); + + for (outer_t j = 0; j < s; ++j) { + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t)*(j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t)*j); + const outer_t sum = outer_curr + outer_prev; std::memcpy(outer_ptr + sizeof(outer_t)*(j+1), &sum, sizeof(outer_t)); + } + const outer_t outer_last = internal::read_as(outer_ptr + sizeof(outer_t)*s); + outer_t idx2 = outer_last; + + const auto inner_routine = [&](outer_t j) { + const outer_t outer_curr = internal::read_as(outer_ptr + sizeof(outer_t)*(j+1)); + const outer_t outer_prev = internal::read_as(outer_ptr + sizeof(outer_t)*j); + Eigen::Map buffer_j(buffer.data() + outer_prev, outer_curr - outer_prev); + outer_t cidx = (2*B) * sizeof(outer_t); + const auto cal0 = calldata.col(2*j); const auto cal1 = calldata.col(2*j+1); const auto anc0 = ancestries.col(2*j); const auto anc1 = ancestries.col(2*j+1); + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + std::memcpy(buffer_j.data() + sizeof(outer_t)*bi, &cidx, sizeof(outer_t)); + auto *n_chunks_ptr = reinterpret_cast(buffer_j.data() + cidx); cidx += sizeof(inner_t); inner_t n_chunks = 0; + for (inner_t k = 0; k < max_chunks; ++k) { + const outer_t base = k * chunk_size; size_t curr_idx = cidx; + auto *chunk_index = reinterpret_cast(buffer_j.data()+curr_idx); curr_idx += sizeof(inner_t); + auto *chunk_nnz = reinterpret_cast(buffer_j.data()+curr_idx); curr_idx += sizeof(chunk_inner_t); + auto *chunk_begin = reinterpret_cast(buffer_j.data()+curr_idx); + inner_t nnz_chunk = 0; + for (inner_t c = 0; c < chunk_size; ++c) { + const outer_t didx = base + c; if (didx >= n) break; + if (cal0[didx]==1 && anc0[didx]==static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + if (cal1[didx]==1 && anc1[didx]==static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } + } + if (nnz_chunk) { std::memcpy(chunk_index, &k, sizeof(inner_t)); *chunk_nnz = nnz_chunk - 1; cidx = curr_idx; ++n_chunks; } + } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + for (size_t bi = 0; bi < B; ++bi) { + const auto a = sel[bi]; + std::memcpy(buffer_j.data() + sizeof(outer_t)*(B + bi), &cidx, sizeof(outer_t)); + auto *n_chunks_ptr = reinterpret_cast(buffer_j.data() + cidx); cidx += sizeof(inner_t); inner_t n_chunks = 0; + for (inner_t k = 0; k < max_chunks; ++k) { const outer_t base = k * chunk_size; size_t curr_idx = cidx; auto *chunk_index = reinterpret_cast(buffer_j.data()+curr_idx); curr_idx += sizeof(inner_t); auto *chunk_nnz = reinterpret_cast(buffer_j.data()+curr_idx); curr_idx += sizeof(chunk_inner_t); auto *chunk_begin = reinterpret_cast(buffer_j.data()+curr_idx); inner_t nnz_chunk = 0; for (inner_t c = 0; c < chunk_size; ++c) { const outer_t didx = base + c; if (didx >= n) break; if (anc0[didx]==static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } if (anc1[didx]==static_cast(a)) { chunk_begin[nnz_chunk++] = c; curr_idx += sizeof(chunk_inner_t); } } if (nnz_chunk) { std::memcpy(chunk_index, &k, sizeof(inner_t)); *chunk_nnz = nnz_chunk - 1; cidx = curr_idx; ++n_chunks; } } + std::memcpy(n_chunks_ptr, &n_chunks, sizeof(inner_t)); + } + if (cidx != static_cast(buffer_j.size())) { /* no-op safety */ } + }; + util::omp_parallel_for(inner_routine, 0, s, n_threads); + + auto file_ptr = fopen_safe(_filename.c_str(), "wb"); + auto fp = file_ptr.get(); + auto total_bytes = std::fwrite(buffer.data(), sizeof(char), idx2, fp); + if (total_bytes != static_cast(idx2)) { + throw util::adelie_core_error("Could not write the full buffer."); + } + return {total_bytes, {}}; +} + +} // namespace io +} // namespace adelie_core + + diff --git a/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.hpp b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.hpp new file mode 100644 index 00000000..df538bd2 --- /dev/null +++ b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.hpp @@ -0,0 +1,77 @@ +#pragma once +#include +#include +#include +#include +#include + +#ifndef ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +#define ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP \ + template +#endif +#ifndef ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R +#define ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R \ + MatrixNaiveSNPCombineR +#endif + +namespace adelie_core { +namespace matrix { + +template < + class ValueType, + class MmapPtrType=std::unique_ptr>, + class IndexType=Eigen::Index +> +class MatrixNaiveSNPCombineR: public MatrixNaiveBase +{ +public: + using base_t = MatrixNaiveBase; + using typename base_t::value_t; + using typename base_t::vec_value_t; + using typename base_t::vec_index_t; + using typename base_t::colmat_value_t; + using typename base_t::rowmat_value_t; + using typename base_t::sp_mat_value_t; + using string_t = std::string; + using io_t = io::IOSNPCombineR; + +private: + const io_t& _io; // IO handler + const size_t _n_threads; // number of threads + vec_value_t _buff; + + inline value_t _cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights, + size_t n_threads, + Eigen::Ref buff + ) const; + + inline value_t _sq_cmul( + int j, + const Eigen::Ref& weights, + Eigen::Ref buff + ) const; + + inline void _ctmul( + int j, + value_t v, + Eigen::Ref out, + size_t n_threads + ) const; + + auto ancestries() const { return _io.ancestries(); } + +public: + explicit MatrixNaiveSNPCombineR( + const io_t& io, + size_t n_threads + ); + + ADELIE_CORE_MATRIX_NAIVE_PURE_OVERRIDE_DECL + ADELIE_CORE_MATRIX_NAIVE_OVERRIDE_DECL +}; + +} // namespace matrix +} // namespace adelie_core \ No newline at end of file diff --git a/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.ipp b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.ipp new file mode 100644 index 00000000..4a9a7710 --- /dev/null +++ b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_r.ipp @@ -0,0 +1,418 @@ +#pragma once +#include +#include +#include +#include + +namespace adelie_core { +namespace matrix { + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::_cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights, + size_t n_threads, + Eigen::Ref buff +) const +{ + return snp_combine_r_dot( + _io, j, v * weights, n_threads, buff + ); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::_sq_cmul( + int j, + const Eigen::Ref& weights, + Eigen::Ref /*buff*/ +) const +{ + using value_t = typename std::decay_t::Scalar; + + const auto A = _io.ancestries(); + const auto snp = j / (1 + A); + const auto col_in_block = j % (1 + A); + + // Count dosages for each sample + std::vector dosages(_io.rows(), 0); + + if (col_in_block == 0) { + // SNP data column + auto it = _io.begin(snp, 0); + const auto end = _io.end(snp, 0); + for (; it != end; ++it) { + dosages[*it]++; + } + } else { + // Ancestry column + const auto anc = col_in_block; + auto it = _io.begin(snp, anc); + const auto end = _io.end(snp, anc); + for (; it != end; ++it) { + dosages[*it]++; + } + } + + // Compute weighted sum of squared dosages + value_t sum = 0; + for (size_t i = 0; i < dosages.size(); ++i) { + const auto dosage = dosages[i]; + sum += weights[i] * dosage * dosage; + } + + return sum; +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::_ctmul( + int j, + value_t v, + Eigen::Ref out, + size_t n_threads +) const +{ + return snp_combine_r_axi( + _io, j, v, out, n_threads + ); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::MatrixNaiveSNPCombineR( + const io_t& io, + size_t n_threads +): + _io(io), + _n_threads(n_threads), + _buff(n_threads * (1 + _io.ancestries())) +{ + if (n_threads < 1) { + throw util::adelie_core_error("n_threads must be >= 1."); + } + if (_io.ancestries() < 1) { + throw util::adelie_core_error("Number of ancestries must be >= 1."); + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights +) +{ + base_t::check_cmul(j, v.size(), weights.size(), rows(), cols()); + return _cmul(j, v, weights, _n_threads, _buff); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::cmul_safe( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights +) const +{ + base_t::check_cmul(j, v.size(), weights.size(), rows(), cols()); + vec_value_t buff(_n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + return _cmul(j, v, weights, _n_threads, buff); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::ctmul( + int j, + value_t v, + Eigen::Ref out +) +{ + base_t::check_ctmul(j, out.size(), rows(), cols()); + _ctmul(j, v, out, _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::bmul( + int j, int q, + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) +{ + base_t::check_bmul(j, q, v.size(), weights.size(), out.size(), rows(), cols()); + if (static_cast(_buff.size()) < q * _n_threads) _buff.resize(q * _n_threads); + snp_combine_r_block_dot( + _io, j, q, v * weights, out, _n_threads, _buff + ); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::bmul_safe( + int j, int q, + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + base_t::check_bmul(j, q, v.size(), weights.size(), out.size(), rows(), cols()); + vec_value_t buff(q * _n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + snp_combine_r_block_dot( + _io, j, q, v * weights, out, _n_threads, buff + ); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::btmul( + int j, int q, + const Eigen::Ref& v, + Eigen::Ref out +) +{ + base_t::check_btmul(j, q, v.size(), out.size(), rows(), cols()); + snp_combine_r_block_axi( + _io, j, q, v, out, _n_threads + ); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::mul( + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + const auto routine = [&](int t) { + out[t] = _cmul(t, v, weights, 1, out /* unused */); + }; + util::omp_parallel_for(routine, 0, cols(), _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::cov( + int j, int q, + const Eigen::Ref& sqrt_weights, + Eigen::Ref out +) const +{ + base_t::check_cov( + j, q, sqrt_weights.size(), + out.rows(), out.cols(), + rows(), cols() + ); + + const auto A = ancestries(); + + out.setZero(); // don't parallelize! q is usually small + + util::rowvec_type bbuff(_io.rows()); + vec_index_t ibuff(_io.rows()); + vec_value_t buff(_n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + bbuff.setZero(); + + int n_solved0 = 0; + while (n_solved0 < q) { + const auto begin0 = j + n_solved0; + const auto snp0 = begin0 / (1 + A); + const auto col_lower0 = begin0 % (1 + A); + const auto col_upper0 = std::min(col_lower0 + q - n_solved0, 1 + A); + + int n_solved1 = 0; + while (n_solved1 <= n_solved0) { + const auto begin1 = j + n_solved1; + const auto col_lower1 = begin1 % (1 + A); + const auto col_upper1 = std::min(col_lower1 + q - n_solved1, 1 + A); + + #if 0 + // This block of code will not be compiled. + if (n_solved0 == n_solved1) { + const auto snp = snp0; + const auto c_low = col_lower0; + const auto c_high = col_upper0; + const auto c_size = c_high - c_low; + + // increase buffer including cross-term computation part as well + if ( + static_cast(buff.size()) < c_size * _n_threads && + _n_threads > 1 && + !util::omp_in_parallel() + ) { + buff.resize(c_size * _n_threads); + } + + const auto routine = [&](auto k0, auto k1) { + const auto sum = snp_combine_r_cross_dot( + _io, + snp * (1 + A) + c_low + k0, + snp * (1 + A) + c_low + k1, + sqrt_weights.square() + ); + out(n_solved0 + k0, n_solved0 + k1) = sum; + }; + + if (_n_threads <= 1 || util::omp_in_parallel()) { + for (int k0 = 0; k0 < static_cast(c_size); ++k0) { + for (int k1 = 0; k1 < static_cast(c_size); ++k1) { + routine(k0, k1); + } + } + } else { + #pragma omp parallel for schedule(static) num_threads(_n_threads) collapse(2) + for (int k0 = 0; k0 < static_cast(c_size); ++k0) { + for (int k1 = 0; k1 < static_cast(c_size); ++k1) { + routine(k0, k1); + } + } + } + + n_solved1 += col_upper1 - col_lower1; + continue; + } + #endif + + /* general routine */ + + const auto col_size0 = col_upper0 - col_lower0; + const auto col_size1 = col_upper1 - col_lower1; + for (size_t c0 = 0; c0 < col_size0; ++c0) { + const auto col_in_block0 = col_lower0 + c0; + size_t nnz = 0; + + if (col_in_block0 == 0) { + // SNP data column + auto it = _io.begin(snp0, 0); + const auto end = _io.end(snp0, 0); + for (; it != end; ++it) { + const auto idx = *it; + if (!bbuff[idx]) { + ibuff[nnz] = idx; + ++nnz; + } + bbuff[idx] += 1; + } + } else { + // Ancestry column + auto it = _io.begin(snp0, col_in_block0); + const auto end = _io.end(snp0, col_in_block0); + for (; it != end; ++it) { + const auto idx = *it; + if (!bbuff[idx]) { + ibuff[nnz] = idx; + ++nnz; + } + bbuff[idx] += 1; + } + } + + for (size_t c1 = 0; c1 < col_size1; ++c1) { + const auto sum = snp_combine_r_dot( + _io, begin1 + c1, + vec_value_t::NullaryExpr(sqrt_weights.size(), [&](auto i) { + const auto sqrt_wi = sqrt_weights[i]; + return sqrt_wi * sqrt_wi * bbuff[i]; + }), + _n_threads, + buff + ); + const auto kk0 = n_solved0 + c0; + const auto kk1 = n_solved1 + c1; + // --- assign symmetrically, no accumulation --- + if (kk0 == kk1) { + // diagonal + out(kk0, kk0) = sum; + } else { + // off–diagonal + out(kk0, kk1) = sum; + out(kk1, kk0) = sum; + } + } + + for (size_t i = 0; i < nnz; ++i) { + bbuff[ibuff[i]] = 0; + } + } + + n_solved1 += col_upper1 - col_lower1; + } + n_solved0 += col_upper0 - col_lower0; + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +int +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::rows() const +{ + return _io.rows(); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +int +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::cols() const +{ + return _io.snps() * (1 + ancestries()); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::sq_mul( + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + const auto routine = [&](int t) { + out[t] = _sq_cmul(t, weights, out /* unused */); + }; + util::omp_parallel_for(routine, 0, cols(), _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::sp_tmul( + const sp_mat_value_t& v, + Eigen::Ref out +) const +{ + base_t::check_sp_tmul( + v.rows(), v.cols(), out.rows(), out.cols(), rows(), cols() + ); + const auto routine = [&](int k) { + typename sp_mat_value_t::InnerIterator it(v, k); + auto out_k = out.row(k); + out_k.setZero(); + for (; it; ++it) { + _ctmul(it.index(), it.value(), out_k, 1); + } + }; + util::omp_parallel_for(routine, 0, v.outerSize(), _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::mean( + const Eigen::Ref&, + Eigen::Ref out +) const +{ + out.setZero(); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_R::var( + const Eigen::Ref&, + const Eigen::Ref&, + Eigen::Ref out +) const +{ + out.setOnes(); +} + +} // namespace matrix +} // namespace adelie_core \ No newline at end of file diff --git a/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.hpp b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.hpp new file mode 100644 index 00000000..861af403 --- /dev/null +++ b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.hpp @@ -0,0 +1,79 @@ +#pragma once +#include +#include +#include +#include +#include + +#ifndef ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +#define ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP \ + template +#endif +#ifndef ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S +#define ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S \ + MatrixNaiveSNPCombineS +#endif + +namespace adelie_core { +namespace matrix { + +template < + class ValueType, + class MmapPtrType=std::unique_ptr>, + class IndexType=Eigen::Index +> +class MatrixNaiveSNPCombineS: public MatrixNaiveBase +{ +public: + using base_t = MatrixNaiveBase; + using typename base_t::value_t; + using typename base_t::vec_value_t; + using typename base_t::vec_index_t; + using typename base_t::colmat_value_t; + using typename base_t::rowmat_value_t; + using typename base_t::sp_mat_value_t; + using string_t = std::string; + using io_t = io::IOSNPCombineS; + +private: + const io_t& _io; // IO handler + const size_t _n_threads; // number of threads + vec_value_t _buff; + + inline value_t _cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights, + size_t n_threads, + Eigen::Ref buff + ) const; + + inline value_t _sq_cmul( + int j, + const Eigen::Ref& weights, + Eigen::Ref buff + ) const; + + inline void _ctmul( + int j, + value_t v, + Eigen::Ref out, + size_t n_threads + ) const; + + auto ancestries() const { return _io.ancestries(); } + +public: + explicit MatrixNaiveSNPCombineS( + const io_t& io, + size_t n_threads + ); + + ADELIE_CORE_MATRIX_NAIVE_PURE_OVERRIDE_DECL + ADELIE_CORE_MATRIX_NAIVE_OVERRIDE_DECL +}; + +} // namespace matrix +} // namespace adelie_core + + diff --git a/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.ipp b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.ipp new file mode 100644 index 00000000..9290cf1d --- /dev/null +++ b/adelie/src/include/adelie_core/matrix/matrix_naive_snp_combine_s.ipp @@ -0,0 +1,423 @@ +#pragma once +#include +#include +#include +#include +#include + +namespace adelie_core { +namespace matrix { + +namespace { + +template +auto both_dot(const IO& io, int j, const V& v, size_t n_threads, Buff& buff) +{ + using value_t = typename std::decay_t::Scalar; + const auto A = io.ancestries(); + const auto snp = j / (2 * A); + const auto col_in_block = j % (2 * A); + const auto nnz = io.nnz()[j]; + const size_t n_bytes = (8 * sizeof(value_t)) * nnz; + if (n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + value_t sum = 0; + auto it = io.begin(snp, col_in_block); + const auto end = io.end(snp, col_in_block); + for (; it != end; ++it) sum += v[*it]; + return sum; + } + auto vbuff = buff.head(n_threads); + vbuff.setZero(); + const size_t n_chunks = io.n_chunks(snp, col_in_block); + const int n_blocks = std::min(n_threads, static_cast(n_chunks)); + if (n_blocks > 0) { + const int block_size = static_cast(n_chunks) / n_blocks; + const int remainder = static_cast(n_chunks) % n_blocks; + #pragma omp parallel for schedule(static) num_threads(n_threads) + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t - remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, col_in_block, begin); + const auto end = io.begin(snp, col_in_block, begin + size); + value_t sum = 0; + for (; it != end; ++it) sum += v[*it]; + vbuff[t] = sum; + } + } + return vbuff.sum(); +} + +template +void both_axi(const IO& io, int j, VT v, Out& out, size_t n_threads) +{ + using value_t = VT; + const auto A = io.ancestries(); + const auto snp = j / (2 * A); + const auto col_in_block = j % (2 * A); + const auto nnz = io.nnz()[j]; + const size_t n_bytes = (4 * sizeof(value_t)) * nnz; + if (n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + auto it = io.begin(snp, col_in_block); + const auto end = io.end(snp, col_in_block); + for (; it != end; ++it) out[*it] += v; + return; + } + const size_t n_chunks = io.n_chunks(snp, col_in_block); + const int n_blocks = std::min(n_threads, static_cast(n_chunks)); + if (n_blocks > 0) { + const int block_size = static_cast(n_chunks) / n_blocks; + const int remainder = static_cast(n_chunks) % n_blocks; + #pragma omp parallel for schedule(static) num_threads(n_threads) + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t - remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, col_in_block, begin); + const auto end = io.begin(snp, col_in_block, begin + size); + for (; it != end; ++it) out[*it] += v; + } + } +} + +} // anonymous namespace + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::_cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights, + size_t n_threads, + Eigen::Ref buff +) const +{ + return both_dot(_io, j, v * weights, n_threads, buff); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::_sq_cmul( + int j, + const Eigen::Ref& weights, + Eigen::Ref /*buff*/ +) const +{ + using value_t = typename std::decay_t::Scalar; + + const auto A = _io.ancestries(); + const auto snp = j / (2 * A); + const auto col_in_block = j % (2 * A); + + // Count dosages for each sample + std::vector dosages(_io.rows(), 0); + + // Both types of columns are simple presence counts + auto it = _io.begin(snp, col_in_block); + const auto end = _io.end(snp, col_in_block); + for (; it != end; ++it) { + dosages[*it]++; + } + + value_t sum = 0; + for (size_t i = 0; i < dosages.size(); ++i) { + const auto dosage = dosages[i]; + sum += weights[i] * dosage * dosage; + } + return sum; +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::_ctmul( + int j, + value_t v, + Eigen::Ref out, + size_t n_threads +) const +{ + return both_axi(_io, j, v, out, n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::MatrixNaiveSNPCombineS( + const io_t& io, + size_t n_threads +): + _io(io), + _n_threads(n_threads), + _buff(n_threads * (2 * _io.ancestries())) +{ + if (n_threads < 1) { + throw util::adelie_core_error("n_threads must be >= 1."); + } + if (_io.ancestries() < 1) { + throw util::adelie_core_error("Number of ancestries must be >= 1."); + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::cmul( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights +) +{ + base_t::check_cmul(j, v.size(), weights.size(), rows(), cols()); + return _cmul(j, v, weights, _n_threads, _buff); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +typename ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::value_t +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::cmul_safe( + int j, + const Eigen::Ref& v, + const Eigen::Ref& weights +) const +{ + base_t::check_cmul(j, v.size(), weights.size(), rows(), cols()); + vec_value_t buff(_n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + return _cmul(j, v, weights, _n_threads, buff); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::ctmul( + int j, + value_t v, + Eigen::Ref out +) +{ + base_t::check_ctmul(j, out.size(), rows(), cols()); + _ctmul(j, v, out, _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::bmul( + int j, int q, + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) +{ + base_t::check_bmul(j, q, v.size(), weights.size(), out.size(), rows(), cols()); + if (static_cast(_buff.size()) < q * _n_threads) _buff.resize(q * _n_threads); + // Compute column-wise dot products independently + for (int k = 0; k < q; ++k) { + out[k] = both_dot(_io, j + k, v * weights, _n_threads, _buff); + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::bmul_safe( + int j, int q, + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + base_t::check_bmul(j, q, v.size(), weights.size(), out.size(), rows(), cols()); + vec_value_t buff(q * _n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + for (int k = 0; k < q; ++k) { + out[k] = both_dot(_io, j + k, v * weights, _n_threads, buff); + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::btmul( + int j, int q, + const Eigen::Ref& v, + Eigen::Ref out +) +{ + base_t::check_btmul(j, q, v.size(), out.size(), rows(), cols()); + for (int k = 0; k < q; ++k) { + both_axi(_io, j + k, v[k], out, _n_threads); + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::mul( + const Eigen::Ref& v, + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + const auto routine = [&](int t) { + out[t] = _cmul(t, v, weights, 1, out /* unused */); + }; + util::omp_parallel_for(routine, 0, cols(), _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::cov( + int j, int q, + const Eigen::Ref& sqrt_weights, + Eigen::Ref out +) const +{ + base_t::check_cov( + j, q, sqrt_weights.size(), + out.rows(), out.cols(), + rows(), cols() + ); + out.setZero(); + + const auto A = ancestries(); + + // Temporary buffers, reused across columns + util::rowvec_type bbuff(_io.rows()); + vec_index_t ibuff(_io.rows()); + vec_value_t buff(_n_threads * (_n_threads > 1) * !util::omp_in_parallel()); + bbuff.setZero(); + + int n_solved0 = 0; + while (n_solved0 < q) { + const auto begin0 = j + n_solved0; + const auto snp0 = begin0 / (2 * A); + const auto col_lower0 = begin0 % (2 * A); + const auto col_upper0 = std::min(col_lower0 + q - n_solved0, 2 * A); + + int n_solved1 = 0; + while (n_solved1 <= n_solved0) { + const auto begin1 = j + n_solved1; + const auto col_lower1 = begin1 % (2 * A); + const auto col_upper1 = std::min(col_lower1 + q - n_solved1, 2 * A); + + const auto col_size0 = col_upper0 - col_lower0; + const auto col_size1 = col_upper1 - col_lower1; + + for (size_t c0 = 0; c0 < static_cast(col_size0); ++c0) { + const auto col_in_block0 = col_lower0 + static_cast(c0); + size_t nnz = 0; + + // Build per-row multiplicity for column 0 within this block + auto it = _io.begin(snp0, col_in_block0); + const auto end = _io.end(snp0, col_in_block0); + for (; it != end; ++it) { + const auto idx = *it; + if (!bbuff[idx]) { + ibuff[nnz] = static_cast(idx); + ++nnz; + } + bbuff[idx] += 1; + } + + // Cross terms with columns in the second block window + for (size_t c1 = 0; c1 < static_cast(col_size1); ++c1) { + // weighted dot using v[i] = w[i] * bbuff[i] + const auto sum = both_dot( + _io, + begin1 + static_cast(c1), + vec_value_t::NullaryExpr(sqrt_weights.size(), [&](auto i) { + const auto sw = sqrt_weights[i]; + return sw * sw * static_cast(bbuff[i]); + }), + _n_threads, + buff + ); + const auto kk0 = n_solved0 + static_cast(c0); + const auto kk1 = n_solved1 + static_cast(c1); + if (kk0 == kk1) { + out(kk0, kk0) = sum; + } else { + out(kk0, kk1) = sum; + out(kk1, kk0) = sum; + } + } + + // reset bbuff entries we touched + for (size_t i = 0; i < nnz; ++i) { + bbuff[ibuff[i]] = 0; + } + } + + n_solved1 += col_upper1 - col_lower1; + } + n_solved0 += col_upper0 - col_lower0; + } +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::mean( + const Eigen::Ref&, + Eigen::Ref out +) const +{ + out.setZero(); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::var( + const Eigen::Ref&, + const Eigen::Ref&, + Eigen::Ref out +) const +{ + out.setOnes(); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +int +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::rows() const +{ + return _io.rows(); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +int +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::cols() const +{ + return _io.snps() * (2 * ancestries()); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::sq_mul( + const Eigen::Ref& weights, + Eigen::Ref out +) const +{ + const auto routine = [&](int t) { + out[t] = _sq_cmul(t, weights, out /* unused */); + }; + util::omp_parallel_for(routine, 0, cols(), _n_threads); +} + +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S_TP +void +ADELIE_CORE_MATRIX_NAIVE_SNP_COMBINE_S::sp_tmul( + const sp_mat_value_t& v, + Eigen::Ref out +) const +{ + base_t::check_sp_tmul( + v.rows(), v.cols(), out.rows(), out.cols(), rows(), cols() + ); + const auto routine = [&](int k) { + typename sp_mat_value_t::InnerIterator it(v, k); + auto out_k = out.row(k); + out_k.setZero(); + for (; it; ++it) { + _ctmul(it.index(), it.value(), out_k, 1); + } + }; + util::omp_parallel_for(routine, 0, v.outerSize(), _n_threads); +} + +} // namespace matrix +} // namespace adelie_core + + diff --git a/adelie/src/include/adelie_core/matrix/utils.hpp b/adelie/src/include/adelie_core/matrix/utils.hpp index 2d68c84b..e2159942 100644 --- a/adelie/src/include/adelie_core/matrix/utils.hpp +++ b/adelie/src/include/adelie_core/matrix/utils.hpp @@ -1,5 +1,7 @@ #pragma once +#include #include +#include #include #include #include @@ -985,6 +987,437 @@ void snp_phased_ancestry_block_axi( } } +template +auto snp_combine_r_dot( + const IOType& io, + int j, + const VType& v, + size_t n_threads, + BuffType& buff +) +{ + using value_t = typename std::decay_t::Scalar; + + const auto A = io.ancestries(); + const auto snp = j / (1 + A); + const auto col_in_block = j % (1 + A); + const auto nnz = io.nnz()[j]; + + // NOTE: multiplier from experimentation + const size_t n_bytes = (8 * sizeof(value_t)) * nnz; + if (col_in_block != 0 || n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + value_t sum = 0; + if (col_in_block == 0) { + // SNP data column + auto it = io.begin(snp, 0); + const auto end = io.end(snp, 0); + for (; it != end; ++it) { + const auto idx = *it; + sum += v[idx]; + } + } else { + // Ancestry column + const auto anc = col_in_block; + auto it = io.begin(snp, anc); + const auto end = io.end(snp, anc); + for (; it != end; ++it) { + const auto idx = *it; + sum += v[idx]; + } + } + return sum; + } + + auto vbuff = buff.head(n_threads); + vbuff.setZero(); + + #pragma omp parallel num_threads(n_threads) + { + if (col_in_block == 0) { + // SNP data column + const size_t n_chunks = io.n_chunks(snp, 0); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, 0, begin); + const auto end = io.begin(snp, 0, begin + size); + + value_t sum = 0; + for (; it != end; ++it) { + const auto idx = *it; + sum += v[idx]; + } + vbuff[t] += sum; + } + } + } else { + // Ancestry column + const auto anc = col_in_block; + const size_t n_chunks = io.n_chunks(snp, anc); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, anc, begin); + const auto end = io.begin(snp, anc, begin + size); + + value_t sum = 0; + for (; it != end; ++it) { + const auto idx = *it; + sum += v[idx]; + } + vbuff[t] += sum; + } + } + } + } + return vbuff.sum(); +} + +template +auto snp_combine_r_cross_dot( + const IOType& io, + int j0, + int j1, + const VType& v +) +{ + using value_t = typename std::decay_t::Scalar; + + const auto A = io.ancestries(); + const auto snp0 = j0 / (1 + A); + const auto col0 = j0 % (1 + A); + const auto snp1 = j1 / (1 + A); + const auto col1 = j1 % (1 + A); + + // Convert column indices to ancestry indices (0 for SNP data) + const auto anc0 = (col0 == 0) ? 0 : col0; + const auto anc1 = (col1 == 0) ? 0 : col1; + + auto it0 = io.begin(snp0, anc0); + const auto end0 = io.end(snp0, anc0); + auto it1 = io.begin(snp1, anc1); + const auto end1 = io.end(snp1, anc1); + + value_t sum = 0; + while ( + (it0 != end0) && + (it1 != end1) + ) { + const auto idx0 = *it0; + const auto idx1 = *it1; + if (idx0 < idx1) { + ++it0; + continue; + } + else if (idx0 > idx1) { + ++it1; + continue; + } + else { + sum += v[idx0]; + ++it0; + ++it1; + } + } + return sum; +} + +template +void snp_combine_r_axi( + const IOType& io, + int j, + ValueType v, + OutType& out, + size_t n_threads +) +{ + using value_t = ValueType; + + const auto A = io.ancestries(); + const auto snp = j / (1 + A); + const auto col_in_block = j % (1 + A); + const auto nnz = io.nnz()[j]; + // NOTE: multiplier from experimentation + const size_t n_bytes = (4 * sizeof(value_t)) * nnz; + if (col_in_block != 0 || n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + if (col_in_block == 0) { + // SNP data column + auto it = io.begin(snp, 0); + const auto end = io.end(snp, 0); + for (; it != end; ++it) { + out[*it] += v; + } + } else { + // Ancestry column + const auto anc = col_in_block; + auto it = io.begin(snp, anc); + const auto end = io.end(snp, anc); + for (; it != end; ++it) { + out[*it] += v; + } + } + return; + } + + if (col_in_block == 0) { + // SNP data column + const size_t n_chunks = io.n_chunks(snp, 0); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp parallel for schedule(static) num_threads(n_threads) + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, 0, begin); + const auto end = io.begin(snp, 0, begin + size); + + for (; it != end; ++it) { + out[*it] += v; + } + } + } + } else { + // Ancestry column + const auto anc = col_in_block; + const size_t n_chunks = io.n_chunks(snp, anc); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp parallel for schedule(static) num_threads(n_threads) + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, anc, begin); + const auto end = io.begin(snp, anc, begin + size); + + for (; it != end; ++it) { + out[*it] += v; + } + } + } + } +} + +template +void snp_combine_r_block_dot( + const IOType& io, + int j, + int q, + const VType& v, + OutType& out, + size_t n_threads, + BuffType& buff +) +{ + using value_t = typename std::decay_t::Scalar; + using rowarr_value_t = util::rowarr_type; + + const auto nnz = io.nnz().segment(j, q).sum(); + // NOTE: multiplier from experimentation + const size_t n_bytes = (8 * sizeof(value_t)) * nnz; + if (n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + for (int k = 0; k < q; ++k) { + out[k] = snp_combine_r_dot(io, j+k, v, n_threads, buff); + } + return; + } + + Eigen::Map mbuff( + buff.data(), q, n_threads + ); + mbuff.setZero(); + + const auto A = io.ancestries(); + + #pragma omp parallel num_threads(n_threads) + { + for (int k = 0; k < q; ++k) { + const auto jj = j + k; + const auto snp = jj / (1 + A); + const auto col_in_block = jj % (1 + A); + + if (col_in_block == 0) { + // SNP data column + const size_t n_chunks = io.n_chunks(snp, 0); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, 0, begin); + const auto end = io.begin(snp, 0, begin + size); + + value_t sum = 0; + for (; it != end; ++it) { + sum += v[*it]; + } + mbuff(k, t) += sum; + } + } + } else { + // Ancestry column + const auto anc = col_in_block; + const size_t n_chunks = io.n_chunks(snp, anc); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, anc, begin); + const auto end = io.begin(snp, anc, begin + size); + + value_t sum = 0; + for (; it != end; ++it) { + sum += v[*it]; + } + mbuff(k, t) += sum; + } + } + } + } + } + + out = mbuff.rowwise().sum(); +} + +template +void snp_combine_r_block_axi( + const IOType& io, + int j, + int q, + const VType& v, + OutType& out, + size_t n_threads +) +{ + using value_t = typename std::decay_t::Scalar; + + const auto nnz = io.nnz().segment(j, q).sum(); + // NOTE: multiplier from experimentation + const size_t n_bytes = (4 * sizeof(value_t)) * nnz; + if (n_threads <= 1 || util::omp_in_parallel() || n_bytes <= Configs::min_bytes) { + for (int k = 0; k < q; ++k) { + snp_combine_r_axi( + io, j+k, v[k], out, n_threads + ); + } + return; + } + + const auto A = io.ancestries(); + + int n_processed = 0; + while (n_processed < q) { + const auto begin = j + n_processed; + const auto snp = begin / (1 + A); + const auto col_lower = begin % (1 + A); + const auto col_upper = std::min(col_lower + q - n_processed, 1 + A); + const auto size = col_upper - col_lower; + + #pragma omp parallel num_threads(n_threads) + { + for (size_t k = 0; k < size; ++k) { + const auto col_in_block = col_lower + k; + const auto vk = v[n_processed + k]; + + if (col_in_block == 0) { + // SNP data column + const size_t n_chunks = io.n_chunks(snp, 0); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, 0, begin); + const auto end = io.begin(snp, 0, begin + size); + + for (; it != end; ++it) { + out[*it] += vk; + } + } + } + } else { + // Ancestry column + const auto anc = col_in_block; + const size_t n_chunks = io.n_chunks(snp, anc); + const int n_blocks = std::min(n_threads, n_chunks); + if (n_blocks > 0) { + const int block_size = n_chunks / n_blocks; + const int remainder = n_chunks % n_blocks; + + #pragma omp for schedule(static) nowait + for (int t = 0; t < n_blocks; ++t) { + const auto begin = ( + std::min(t, remainder) * (block_size + 1) + + std::max(t-remainder, 0) * block_size + ); + const auto size = block_size + (t < remainder); + auto it = io.begin(snp, anc, begin); + const auto end = io.begin(snp, anc, begin + size); + + for (; it != end; ++it) { + out[*it] += vk; + } + } + } + } + } + } + + n_processed += size; + } +} + template void dgemtm( const MType& m, @@ -1013,4 +1446,4 @@ void dgemtm( } } // namespace matrix -} // namespace adelie_core \ No newline at end of file +} // namespace adelie_core diff --git a/adelie/src/include/adelie_core/optimization/search_pivot.hpp b/adelie/src/include/adelie_core/optimization/search_pivot.hpp index 760eb5d5..57ad1228 100644 --- a/adelie/src/include/adelie_core/optimization/search_pivot.hpp +++ b/adelie/src/include/adelie_core/optimization/search_pivot.hpp @@ -1,4 +1,5 @@ #pragma once +#include namespace adelie_core { namespace optimization { @@ -63,4 +64,4 @@ int search_pivot( } } // namespace optimization -} // namespace adelie_core \ No newline at end of file +} // namespace adelie_core diff --git a/adelie/src/include/adelie_core/util/macros.hpp b/adelie/src/include/adelie_core/util/macros.hpp index 3bdc476d..b5621e95 100644 --- a/adelie/src/include/adelie_core/util/macros.hpp +++ b/adelie/src/include/adelie_core/util/macros.hpp @@ -1,4 +1,5 @@ #pragma once +#include #ifndef ADELIE_CORE_STRONG_INLINE #if defined(_MSC_VER) diff --git a/adelie/src/py_io.cpp b/adelie/src/py_io.cpp index b431ff4d..8d9c5049 100644 --- a/adelie/src/py_io.cpp +++ b/adelie/src/py_io.cpp @@ -1,5 +1,6 @@ #include "py_decl.hpp" #include +#include namespace py = pybind11; namespace ad = adelie_core; @@ -183,9 +184,189 @@ void io_snp_phased_ancestry(py::module_& m) ; } +void io_snp_combine_r(py::module_& m) +{ + using io_t = ad::io::IOSNPCombineR<>; + using base_t = typename io_t::base_t; + using string_t = typename io_t::string_t; + using colarr_value_t = typename io_t::colarr_value_t; + using uint_vec_t = std::vector; + py::class_(m, "IOSNPCombineR") + .def(py::init< + const string_t&, + const string_t& + >(), + py::arg("filename"), + py::arg("read_mode") + ) + .def_property_readonly("rows", &io_t::rows, "Number of rows.") + .def_property_readonly("snps", &io_t::snps, "Number of SNPs.") + .def_property_readonly("cols", &io_t::cols, "Number of columns.") + .def_property_readonly("ancestries", &io_t::ancestries, "Number of ancestries.") + .def_property_readonly("nnz", &io_t::nnz, "Number of non-zero entries for each column.") + .def("to_dense", &io_t::to_dense, + py::arg("n_threads")=1, + R"delimiter( + Creates a dense SNP unphased, ancestry matrix from the file. + + Parameters + ---------- + n_threads : int, optional + Number of threads. + Default is ``1``. + + Returns + ------- + dense : (n, s*(1+A)) ndarray + Dense SNP unphased, ancestry matrix. + For each SNP j, the columns are: + - Column j*(1+A): SNP data + - Columns j*(1+A)+1 to j*(1+A)+A: Ancestry dosages + )delimiter") + .def("write", []( + const io_t& io, + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads + ) { + std::tuple> out; + std::string error; + try { + out = io.write(calldata, ancestries, A, n_threads); + } catch (const std::exception& e) { + error = e.what(); + } + return std::make_tuple( + std::get<0>(out), + std::get<1>(out), + error + ); + }, + py::arg("calldata").noconvert(), + py::arg("ancestries").noconvert(), + py::arg("A"), + py::arg("n_threads") + ) + .def("write", []( + const io_t& io, + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const uint_vec_t& selected_ancestries, + size_t n_threads + ) { + std::tuple> out; + std::string error; + try { + out = io.write(calldata, ancestries, A_total, selected_ancestries, n_threads); + } catch (const std::exception& e) { + error = e.what(); + } + return std::make_tuple( + std::get<0>(out), + std::get<1>(out), + error + ); + }, + py::arg("calldata").noconvert(), + py::arg("ancestries").noconvert(), + py::arg("A"), + py::arg("selected_ancestries"), + py::arg("n_threads") + ) + ; +} + +void io_snp_combine_s(py::module_& m) +{ + using io_t = ad::io::IOSNPCombineS<>; + using base_t = typename io_t::base_t; + using string_t = typename io_t::string_t; + using colarr_value_t = typename io_t::colarr_value_t; + using uint_vec_t = std::vector; + py::class_(m, "IOSNPCombineS") + .def(py::init< + const string_t&, + const string_t& + >(), + py::arg("filename"), + py::arg("read_mode") + ) + .def_property_readonly("rows", &io_t::rows, "Number of rows.") + .def_property_readonly("snps", &io_t::snps, "Number of SNPs.") + .def_property_readonly("cols", &io_t::cols, "Number of columns.") + .def_property_readonly("ancestries", &io_t::ancestries, "Number of ancestries.") + .def_property_readonly("nnz", &io_t::nnz, "Number of non-zero entries for each column.") + .def("to_dense", &io_t::to_dense, + py::arg("n_threads")=1, + R"delimiter( + Creates a dense SNP both-ancestry matrix from the file. + + Returns a matrix of shape (n, s*(2*A)) where for each SNP j: + - Columns j*(2*A) .. j*(2*A)+(A-1): mutated haplotype counts per ancestry (0..2) + - Columns j*(2*A)+A .. j*(2*A)+(2*A-1): ancestry dosage counts per ancestry (0..2) + )delimiter") + .def("write", []( + const io_t& io, + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A, + size_t n_threads + ) { + std::tuple> out; + std::string error; + try { + out = io.write(calldata, ancestries, A, n_threads); + } catch (const std::exception& e) { + error = e.what(); + } + return std::make_tuple( + std::get<0>(out), + std::get<1>(out), + error + ); + }, + py::arg("calldata").noconvert(), + py::arg("ancestries").noconvert(), + py::arg("A"), + py::arg("n_threads") + ) + .def("write", []( + const io_t& io, + const Eigen::Ref& calldata, + const Eigen::Ref& ancestries, + size_t A_total, + const uint_vec_t& selected_ancestries, + size_t n_threads + ) { + std::tuple> out; + std::string error; + try { + out = io.write(calldata, ancestries, A_total, selected_ancestries, n_threads); + } catch (const std::exception& e) { + error = e.what(); + } + return std::make_tuple( + std::get<0>(out), + std::get<1>(out), + error + ); + }, + py::arg("calldata").noconvert(), + py::arg("ancestries").noconvert(), + py::arg("A"), + py::arg("selected_ancestries"), + py::arg("n_threads") + ) + ; +} + void register_io(py::module_& m) { io_snp_base(m); io_snp_unphased(m); io_snp_phased_ancestry(m); + io_snp_combine_r(m); + io_snp_combine_s(m); } \ No newline at end of file diff --git a/adelie/src/py_matrix.cpp b/adelie/src/py_matrix.cpp index dbbf62a8..6bca4a3a 100644 --- a/adelie/src/py_matrix.cpp +++ b/adelie/src/py_matrix.cpp @@ -1661,6 +1661,74 @@ void matrix_naive_snp_phased_ancestry(py::module_& m, const char* name) ; } +template +void matrix_naive_snp_combine_r(py::module_& m, const char* name) +{ + using internal_t = ad::matrix::MatrixNaiveSNPCombineR; + using base_t = typename internal_t::base_t; + using io_t = typename internal_t::io_t; + py::class_(m, name, + "Core matrix class for naive SNP unphased, ancestry matrix." + ) + .def( + py::init< + const io_t&, + size_t + >(), + py::arg("io"), + py::arg("n_threads") + ) + .def("mean", &internal_t::mean, R"delimiter( + Computes the implied column means. + + The implied column means are zero. + + Parameters + ---------- + weights : (n,) ndarray + Vector of weights. + out : (p,) ndarray + Vector to store in-place the result. + )delimiter") + .def("var", &internal_t::var, R"delimiter( + Computes the implied column variances. + + The implied column variances are one. + + Parameters + ---------- + centers : (p,) ndarray + Vector of centers. + weights : (n,) ndarray + Vector of weights. + out : (p,) ndarray + Vector to store in-place the result. + )delimiter") + ; +} + +template +void matrix_naive_snp_combine_s(py::module_& m, const char* name) +{ + using internal_t = ad::matrix::MatrixNaiveSNPCombineS; + using base_t = typename internal_t::base_t; + using io_t = typename internal_t::io_t; + py::class_(m, name, + "Core matrix class for naive SNP both-ancestry matrix." + ) + .def( + py::init< + const io_t&, + size_t + >(), + py::arg("io"), + py::arg("n_threads") + ) + .def("mean", &internal_t::mean) + .def("var", &internal_t::var) + ; +} + template void matrix_naive_sparse(py::module_& m, const char* name) { @@ -1963,6 +2031,10 @@ void register_matrix(py::module_& m) matrix_naive_snp_unphased(m, "MatrixNaiveSNPUnphased32"); matrix_naive_snp_phased_ancestry(m, "MatrixNaiveSNPPhasedAncestry64"); matrix_naive_snp_phased_ancestry(m, "MatrixNaiveSNPPhasedAncestry32"); + matrix_naive_snp_combine_r(m, "MatrixNaiveSNPCombineR64"); + matrix_naive_snp_combine_r(m, "MatrixNaiveSNPCombineR32"); + matrix_naive_snp_combine_s(m, "MatrixNaiveSNPCombineS64"); + matrix_naive_snp_combine_s(m, "MatrixNaiveSNPCombineS32"); matrix_naive_sparse>(m, "MatrixNaiveSparse64F"); matrix_naive_sparse>(m, "MatrixNaiveSparse32F"); diff --git a/adelie/src/py_matrix_utils.cpp b/adelie/src/py_matrix_utils.cpp index 98ecce72..cf5c75b4 100644 --- a/adelie/src/py_matrix_utils.cpp +++ b/adelie/src/py_matrix_utils.cpp @@ -1,5 +1,6 @@ #include "py_decl.hpp" #include +#include #include #include #include @@ -28,6 +29,7 @@ void utils(py::module_& m) using cref_mvec_value_t = Eigen::Ref>; using snp_unphased_io_t = ad::io::IOSNPUnphased<>; using snp_phased_ancestry_io_t = ad::io::IOSNPPhasedAncestry<>; + using snp_combine_r_io_t = ad::io::IOSNPCombineR<>; using sw_t = ad::util::Stopwatch; m.def("dvaddi", ad::matrix::dvaddi); @@ -345,6 +347,47 @@ void utils(py::module_& m) } return time_elapsed / n_sims; }); + + m.def("bench_snp_combine_r_dot", []( + const snp_combine_r_io_t& io, + int j, + cref_vec_value_t& v, + size_t n_threads, + size_t n_sims + ){ + sw_t sw; + double time_elapsed = 0; + volatile double out = 0; + vec_value_t buff(n_threads); + for (size_t i = 0; i < n_sims; ++i) { + sw.start(); + out += ad::matrix::snp_combine_r_dot( + io, j, v, n_threads, buff + ); + time_elapsed += sw.elapsed(); + } + return time_elapsed / n_sims; + }); + + m.def("bench_snp_combine_r_axi", []( + const snp_combine_r_io_t& io, + int j, + cref_vec_value_t& v, + size_t n_threads, + size_t n_sims + ){ + sw_t sw; + double time_elapsed = 0; + vec_value_t out(v.size()); + for (size_t i = 0; i < n_sims; ++i) { + sw.start(); + ad::matrix::snp_combine_r_axi( + io, j, v[0], out, n_threads + ); + time_elapsed += sw.elapsed(); + } + return time_elapsed / n_sims; + }); } void register_matrix_utils(py::module_& m) diff --git a/adelie/src/src/io/io.hpp b/adelie/src/src/io/io.hpp index 91963a43..739a33ed 100644 --- a/adelie/src/src/io/io.hpp +++ b/adelie/src/src/io/io.hpp @@ -2,7 +2,11 @@ #include #include #include +#include +#include extern template class adelie_core::io::IOSNPBase<>; extern template class adelie_core::io::IOSNPPhasedAncestry<>; -extern template class adelie_core::io::IOSNPUnphased<>; \ No newline at end of file +extern template class adelie_core::io::IOSNPUnphased<>; +extern template class adelie_core::io::IOSNPCombineR<>; +extern template class adelie_core::io::IOSNPCombineS<>; \ No newline at end of file diff --git a/adelie/src/src/io/io_snp_combine_r.cpp b/adelie/src/src/io/io_snp_combine_r.cpp new file mode 100644 index 00000000..9ef1c4ed --- /dev/null +++ b/adelie/src/src/io/io_snp_combine_r.cpp @@ -0,0 +1,4 @@ +#include +#include + +template class adelie_core::io::IOSNPCombineR<>; \ No newline at end of file diff --git a/adelie/src/src/io/io_snp_combine_s.cpp b/adelie/src/src/io/io_snp_combine_s.cpp new file mode 100644 index 00000000..c07c08bc --- /dev/null +++ b/adelie/src/src/io/io_snp_combine_s.cpp @@ -0,0 +1,6 @@ +#include +#include + +template class adelie_core::io::IOSNPCombineS<>; + + diff --git a/adelie/src/src/matrix/matrix.hpp b/adelie/src/src/matrix/matrix.hpp index d3b66585..c57d6fa3 100644 --- a/adelie/src/src/matrix/matrix.hpp +++ b/adelie/src/src/matrix/matrix.hpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include #include @@ -106,6 +108,12 @@ extern template class adelie_core::matrix::MatrixNaiveSNPPhasedAncestry; extern template class adelie_core::matrix::MatrixNaiveSNPUnphased; extern template class adelie_core::matrix::MatrixNaiveSNPUnphased; +extern template class adelie_core::matrix::MatrixNaiveSNPCombineR; +extern template class adelie_core::matrix::MatrixNaiveSNPCombineR; + +extern template class adelie_core::matrix::MatrixNaiveSNPCombineS; +extern template class adelie_core::matrix::MatrixNaiveSNPCombineS; + extern template class adelie_core::matrix::MatrixNaiveSparse>; extern template class adelie_core::matrix::MatrixNaiveSparse>; diff --git a/adelie/src/src/matrix/matrix_naive_snp_combine_r.cpp b/adelie/src/src/matrix/matrix_naive_snp_combine_r.cpp new file mode 100644 index 00000000..b6f1e71e --- /dev/null +++ b/adelie/src/src/matrix/matrix_naive_snp_combine_r.cpp @@ -0,0 +1,5 @@ +#include +#include + +template class adelie_core::matrix::MatrixNaiveSNPCombineR; +template class adelie_core::matrix::MatrixNaiveSNPCombineR; \ No newline at end of file diff --git a/adelie/src/src/matrix/matrix_naive_snp_combine_s.cpp b/adelie/src/src/matrix/matrix_naive_snp_combine_s.cpp new file mode 100644 index 00000000..aefb0b13 --- /dev/null +++ b/adelie/src/src/matrix/matrix_naive_snp_combine_s.cpp @@ -0,0 +1,7 @@ +#include +#include + +template class adelie_core::matrix::MatrixNaiveSNPCombineS; +template class adelie_core::matrix::MatrixNaiveSNPCombineS; + + diff --git a/benchmark/bench_matrix.py b/benchmark/bench_matrix.py index 3f7cdf25..22f2b04d 100644 --- a/benchmark/bench_matrix.py +++ b/benchmark/bench_matrix.py @@ -220,6 +220,8 @@ def bench_io( "snp_unphased_axi": utils.bench_snp_unphased_axi, "snp_phased_ancestry_dot": utils.bench_snp_phased_ancestry_dot, "snp_phased_ancestry_axi": utils.bench_snp_phased_ancestry_axi, + "snp_combine_r_dot": utils.bench_snp_combine_r_dot, + "snp_combine_r_axi": utils.bench_snp_combine_r_axi, }[method] ad.configs.set_configs("min_bytes", min_bytes) @@ -232,10 +234,17 @@ def bench_io( for j, n in enumerate(n_list): filename = "/tmp/bench_snp_tmp.snpdat" if "_unphased_" in method: - io = ad.io.snp_unphased(filename) - data = ad.data.snp_unphased(n, 1, one_ratio=0.35, missing_ratio=0.1, two_ratio=0.05) - io.write(data["X"]) - io.read() + if "_ancestry_" in method: + A = 8 + io = ad.io.snp_combine_r(filename) + data = ad.data.snp_combine_r(n, 1, A, one_ratio=0.45, two_ratio=0.05) + io.write(data["X"], data["ancestries"], A) + io.read() + else: + io = ad.io.snp_unphased(filename) + data = ad.data.snp_unphased(n, 1, one_ratio=0.35, missing_ratio=0.1, two_ratio=0.05) + io.write(data["X"]) + io.read() elif "_phased_" in method: io = ad.io.snp_phased_ancestry(filename) A = 8 diff --git a/demos/toy_combine_r_example.py b/demos/toy_combine_r_example.py new file mode 100644 index 00000000..5a172126 --- /dev/null +++ b/demos/toy_combine_r_example.py @@ -0,0 +1,207 @@ +import numpy as np +import adelie.io as io +import adelie.data as data +import os +import time + +# This example demonstrates how to create a .snpdat file for +# SNP unphased ancestry data, write it to a file, read it back, +# and verify the contents. + +# Set a seed for reproducibility +np.random.seed(0) + +# Toy data parameters +n_samples = 5 +n_snps = 7 +n_ancestries = 3 + +print("Generating toy data with parameters:") +print(f" - Samples: {n_samples}") +print(f" - SNPs: {n_snps}") +print(f" - Ancestries: {n_ancestries}") + +# Generate toy data using snp_combine_r +data_dict = data.snp_combine_r( + n=n_samples, + s=n_snps, + A=n_ancestries, + seed=0 +) + +# Extract calldata and ancestries from the generated data +calldata = data_dict["X"] +ancestries = data_dict["ancestries"] + +print(f"\nCalldata matrix (shape: {calldata.shape}):") +print(calldata) + +print(f"\nAncestries matrix (shape: {ancestries.shape}):") +print(ancestries) + +# Benchmark the nested for loops approach +print("\nBenchmarking nested for loops approach...") +start_time = time.time() + +# For verification, create the expected dense matrix from original data. +# The structure of the dense matrix is explained in the docstring of +# adelie.io.snp_combine_r. +expected_dense = np.zeros((n_samples, n_snps * (1 + n_ancestries)), dtype=np.int8) +for i in range(n_samples): + for j in range(n_snps): + # First column of each SNP block is the calldata value + expected_dense[i, j * (1 + n_ancestries)] = calldata[i, j] + + # Next A columns are ancestry dosages + # For each ancestry type, count how many haplotypes have that ancestry + for a in range(n_ancestries): + # Count ancestries for both haplotypes at this SNP position + if ancestries[i, 2*j] == a: + expected_dense[i, j * (1 + n_ancestries) + 1 + a] += 1 + if ancestries[i, 2*j + 1] == a: + expected_dense[i, j * (1 + n_ancestries) + 1 + a] += 1 + +nested_loops_time = time.time() - start_time +print(f"Nested loops approach took: {nested_loops_time:.6f} seconds") + +# Vectorized approach +print("\nBenchmarking vectorized approach...") +start_time = time.time() + +# Create the dense matrix using vectorized operations +expected_dense_vectorized = np.zeros((n_samples, n_snps * (1 + n_ancestries)), dtype=np.int8) + +# Set the calldata values (first column of each SNP block) +expected_dense_vectorized[:, ::(1 + n_ancestries)] = calldata + +# Create ancestry counts using broadcasting +for a in range(n_ancestries): + # Count ancestries for both haplotypes at each SNP position + ancestry_counts = (ancestries[:, ::2] == a).astype(np.int8) + (ancestries[:, 1::2] == a).astype(np.int8) + # Place counts in the appropriate columns + expected_dense_vectorized[:, a+1::(1 + n_ancestries)] = ancestry_counts + +vectorized_time = time.time() - start_time +print(f"Vectorized approach took: {vectorized_time:.6f} seconds") + +# Verify both approaches produce the same result +np.testing.assert_array_equal(expected_dense, expected_dense_vectorized) +print("Verification successful: vectorized approach matches nested loops approach") + +print(f"\nExpected dense matrix (shape: {expected_dense.shape}):") +print(expected_dense) + +# Define output file path +output_filename = "combine_r.snpdat" + +# Create an IO handler for SNP unphased ancestry data. +handler = io.snp_combine_r(filename=output_filename) + +# Write the data to the file in .snpdat format. +print(f"\nWriting data to '{output_filename}'...") +try: + # Benchmark the io.snp_combine_r approach + print("\nBenchmarking io.snp_combine_r approach...") + start_time = time.time() + total_bytes, benchmark = handler.write( + calldata=calldata, + ancestries=ancestries, + A=n_ancestries, + n_threads=2 # Example of using 2 threads + ) + + print(f"Successfully wrote {total_bytes} bytes.") + print("Benchmark timings for writing:") + for key, value in benchmark.items(): + print(f" - {key}: {value:.6f} seconds") + + # Now, demonstrate reading the data back. + print(f"\nReading data back from '{output_filename}'...") + + # Create a new handler for reading. + handler_read = io.snp_combine_r(filename=output_filename, read_mode="mmap") + + handler_read.read() + + # Verify the properties of the read data. + print("Verifying properties of the read data...") + assert handler_read.is_read + assert handler_read.rows == n_samples + assert handler_read.snps == n_snps + assert handler_read.ancestries == n_ancestries + assert handler_read.cols == n_snps * (1 + n_ancestries) # Each SNP has 1 + n_ancestries columns + print(" - Properties (rows, snps, ancestries, cols) are correct.") + + # Reconstruct the dense matrix from the file. + # The to_dense() method returns a matrix of shape (n_samples, n_snps * (1 + n_ancestries)). + dense_matrix_from_file = handler_read.to_dense(n_threads=2) + io_approach_time = time.time() - start_time + print(f"IO approach took: {io_approach_time:.6f} seconds") + + print(f"\nDense matrix from file (shape: {dense_matrix_from_file.shape}):") + print(dense_matrix_from_file) + + # Compare the reconstructed matrix with the expected one. + np.testing.assert_array_equal(dense_matrix_from_file, expected_dense) + print(" - Verification successful: reconstructed matrix matches expected matrix.") + + # ------------------------------------------------------------------ + # Demonstrate selected ancestries (e.g., only ancestries 0 and 2) + # ------------------------------------------------------------------ + selected = np.array([0, 2], dtype=np.uint32) + output_filename_sel = "combine_r_selected.snpdat" + handler_sel = io.snp_combine_r(filename=output_filename_sel) + # Write only selected ancestries; output has (1 + len(selected)) columns per SNP + total_bytes_sel, benchmark_sel = handler_sel.write( + calldata=calldata, + ancestries=ancestries, + A=n_ancestries, + n_threads=2, + selected_ancestries=selected + ) + + handler_sel_read = io.snp_combine_r(filename=output_filename_sel, read_mode="mmap") + handler_sel_read.read() + + assert handler_sel_read.rows == n_samples + assert handler_sel_read.snps == n_snps + assert handler_sel_read.ancestries == len(selected) + assert handler_sel_read.cols == n_snps * (1 + len(selected)) + + # Build expected dense restricted to selected ancestries + expected_dense_sel = np.zeros((n_samples, n_snps * (1 + len(selected))), dtype=np.int8) + # Genotype column stays identical + expected_dense_sel[:, ::(1 + len(selected))] = calldata + # Ancestry dosage columns only for selected labels, in the same order as 'selected' + for ai, a in enumerate(selected): + ancestry_counts = (ancestries[:, ::2] == a).astype(np.int8) + (ancestries[:, 1::2] == a).astype(np.int8) + expected_dense_sel[:, ai+1::(1 + len(selected))] = ancestry_counts + + dense_matrix_from_file_sel = handler_sel_read.to_dense(n_threads=2) + + # Print both matrices for comparison + print(f"\nExpected dense matrix for selected ancestries {selected} (shape: {expected_dense_sel.shape}):") + print(expected_dense_sel) + + print(f"\nActual dense matrix from file for selected ancestries {selected} (shape: {dense_matrix_from_file_sel.shape}):") + print(dense_matrix_from_file_sel) + + np.testing.assert_array_equal(dense_matrix_from_file_sel, expected_dense_sel) + print(" - Selected ancestries mode verified.") + + # Print performance comparison + print("\nPerformance Comparison:") + print(f"Nested loops approach: {nested_loops_time:.6f} seconds") + print(f"Vectorized approach: {vectorized_time:.6f} seconds") + print(f"IO approach: {io_approach_time:.6f} seconds") + print(f"Speedup factor (vectorized vs nested): {nested_loops_time/vectorized_time:.2f}x") + print(f"Speedup factor (IO vs nested): {nested_loops_time/io_approach_time:.2f}x") + +finally: + # Clean up the created file. + if os.path.exists(output_filename): + os.remove(output_filename) + print(f"\nCleaned up '{output_filename}'.") + if os.path.exists(output_filename_sel): + os.remove(output_filename_sel) + print(f"Cleaned up '{output_filename_sel}'.") diff --git a/demos/toy_combine_s_example.py b/demos/toy_combine_s_example.py new file mode 100644 index 00000000..0c2d9dfb --- /dev/null +++ b/demos/toy_combine_s_example.py @@ -0,0 +1,147 @@ +import numpy as np +import adelie.io as io +import adelie.data as data +import os +import time + +# This example demonstrates how to create a .snpdat file for +# SNP both-ancestry data, write it to a file, read it back, +# and verify the contents, including a selected-ancestries mode. + +np.random.seed(0) + +# Toy data parameters +n_samples = 5 +n_snps = 4 +n_ancestries = 3 + +print("Generating toy phased data with parameters:") +print(f" - Samples: {n_samples}") +print(f" - SNPs: {n_snps}") +print(f" - Ancestries: {n_ancestries}") + +# Generate toy phased calldata (0/1 per hap) and phased ancestries +D = data.snp_phased_ancestry(n=n_samples, s=n_snps, A=n_ancestries, seed=0) +calldata = D["X"] # shape (n, 2*s), entries 0/1 (per hap) +ancestries = D["ancestries"] # shape (n, 2*s), entries in [0, A) + +print(f"\nPhased calldata (shape: {calldata.shape}):") +print(calldata) + +print(f"\nPhased ancestries (shape: {ancestries.shape}):") +print(ancestries) + +# Build expected dense matrix for the both-ancestry format +# Per SNP j, block of size (2*A): +# - First A cols: mutated hap counts per ancestry a +# - Next A cols: ancestry dosage counts per ancestry a + +print("\nComputing expected both-ancestry dense matrix (full A)...") +start = time.time() +expected_dense = np.zeros((n_samples, n_snps * (2 * n_ancestries)), dtype=np.int8) + +# Convenience views per hap +cal0 = calldata[:, ::2] +cal1 = calldata[:, 1::2] +anc0 = ancestries[:, ::2] +anc1 = ancestries[:, 1::2] + +for a in range(n_ancestries): + # mutated hap counts for ancestry a + mut_counts = ((cal0 == 1) & (anc0 == a)).astype(np.int8) + ((cal1 == 1) & (anc1 == a)).astype(np.int8) + expected_dense[:, a::(2 * n_ancestries)] = mut_counts + # ancestry dosage counts for ancestry a + dosage_counts = (anc0 == a).astype(np.int8) + (anc1 == a).astype(np.int8) + expected_dense[:, (n_ancestries + a)::(2 * n_ancestries)] = dosage_counts + +compute_time = time.time() - start +print(f"Done. Time: {compute_time:.6f}s") +print(f"Expected dense (full) shape: {expected_dense.shape}") +print(expected_dense) + +# Write using io.snp_combine_s +output_filename = "combine_s.snpdat" +handler = io.snp_combine_s(filename=output_filename) + +print(f"\nWriting both-ancestry data to '{output_filename}'...") +start = time.time() +total_bytes, benchmark = handler.write( + calldata=calldata, + ancestries=ancestries, + A=n_ancestries, + n_threads=2, +) +write_time = time.time() - start +print(f"Successfully wrote {total_bytes} bytes in {write_time:.6f}s.") + +# Read back and verify +handler_read = io.snp_combine_s(filename=output_filename, read_mode="mmap") +handler_read.read() + +print("Verifying metadata (full A)...") +assert handler_read.rows == n_samples +assert handler_read.snps == n_snps +assert handler_read.ancestries == n_ancestries +assert handler_read.cols == n_snps * (2 * n_ancestries) +print(" - OK") + +print("Reconstructing dense matrix from file (full A)...") +dense_from_file = handler_read.to_dense(n_threads=2) +np.testing.assert_array_equal(dense_from_file, expected_dense) +print(" - Reconstruction matches expected.") + +# ------------------------------------------------------------------ +# Selected ancestries demo (e.g., only ancestries 0 and 2) +# ------------------------------------------------------------------ +selected = np.array([0, 2], dtype=np.uint32) +output_filename_sel = "combine_s_selected.snpdat" +handler_sel = io.snp_combine_s(filename=output_filename_sel) + +print(f"\nWriting both-ancestry data with selected ancestries {selected.tolist()} to '{output_filename_sel}'...") +start = time.time() +total_bytes_sel, benchmark_sel = handler_sel.write( + calldata=calldata, + ancestries=ancestries, + A=n_ancestries, + n_threads=2, + selected_ancestries=selected, +) +print(f"Successfully wrote {total_bytes_sel} bytes in {time.time() - start:.6f}s.") + +handler_sel_read = io.snp_combine_s(filename=output_filename_sel, read_mode="mmap") +handler_sel_read.read() + +print("Verifying metadata (selected ancestries)...") +assert handler_sel_read.rows == n_samples +assert handler_sel_read.snps == n_snps +assert handler_sel_read.ancestries == len(selected) +assert handler_sel_read.cols == n_snps * (2 * len(selected)) +print(" - OK") + +print("Computing expected dense for selected ancestries...") +expected_dense_sel = np.zeros((n_samples, n_snps * (2 * len(selected))), dtype=np.int8) +for ai, a in enumerate(selected): + # mutated counts for selected ancestry a + mut_counts = ((cal0 == 1) & (anc0 == a)).astype(np.int8) + ((cal1 == 1) & (anc1 == a)).astype(np.int8) + expected_dense_sel[:, ai::(2 * len(selected))] = mut_counts + # dosage counts for selected ancestry a + dosage_counts = (anc0 == a).astype(np.int8) + (anc1 == a).astype(np.int8) + expected_dense_sel[:, (len(selected) + ai)::(2 * len(selected))] = dosage_counts + +print(f"Expected dense (selected) shape: {expected_dense_sel.shape}") +print(expected_dense_sel) + +print("Reconstructing dense matrix from file (selected ancestries)...") +dense_from_file_sel = handler_sel_read.to_dense(n_threads=2) +print(f"Actual dense (selected) shape: {dense_from_file_sel.shape}") +print(dense_from_file_sel) +np.testing.assert_array_equal(dense_from_file_sel, expected_dense_sel) +print(" - Selected ancestries mode verified.") + +# Cleanup +for fn in [output_filename, output_filename_sel]: + if os.path.exists(fn): + os.remove(fn) + print(f"Cleaned up '{fn}'.") + + diff --git a/docs/sphinx/api_reference/index.rst b/docs/sphinx/api_reference/index.rst index dae5472e..30375d9f 100644 --- a/docs/sphinx/api_reference/index.rst +++ b/docs/sphinx/api_reference/index.rst @@ -98,6 +98,7 @@ adelie.data dense snp_phased_ancestry snp_unphased + snp_combine_r adelie.diagnostic @@ -168,6 +169,7 @@ adelie.io snp_phased_ancestry snp_unphased + snp_combine_r adelie.matrix @@ -195,6 +197,7 @@ adelie.matrix one_hot snp_phased_ancestry snp_unphased + snp_combine_r sparse standardize subset diff --git a/docs/sphinx/api_reference/internal.rst b/docs/sphinx/api_reference/internal.rst index f855bb56..3e9e81d1 100644 --- a/docs/sphinx/api_reference/internal.rst +++ b/docs/sphinx/api_reference/internal.rst @@ -102,6 +102,10 @@ Internal adelie_core.matrix.MatrixNaiveSNPPhasedAncestry64 adelie_core.matrix.MatrixNaiveSNPUnphased32 adelie_core.matrix.MatrixNaiveSNPUnphased64 + adelie_core.matrix.MatrixNaiveSNPCombineR32 + adelie_core.matrix.MatrixNaiveSNPCombineR64 + adelie_core.matrix.MatrixNaiveSNPCombineS32 + adelie_core.matrix.MatrixNaiveSNPCombineS64 adelie_core.matrix.MatrixNaiveSparse32F adelie_core.matrix.MatrixNaiveSparse64F adelie_core.matrix.MatrixNaiveStandardize32 diff --git a/setup.py b/setup.py index 1c99029e..37e8aa62 100644 --- a/setup.py +++ b/setup.py @@ -25,6 +25,81 @@ def run_cmd(cmd): return output.rstrip() +def maybe_run_cmd(cmd): + try: + return run_cmd(cmd) + except RuntimeError: + return None + + +def append_unique_dirs(dirs, new_dirs): + for new_dir in new_dirs: + if new_dir and os.path.isdir(new_dir) and new_dir not in dirs: + dirs.append(new_dir) + + +def has_eigen_headers(include_dir): + return os.path.isfile(os.path.join(include_dir, "Eigen", "Core")) + + +def resolve_eigen_include_dirs(conda_prefix, system_name): + include_candidates = [] + + for env_var in ["EIGEN3_INCLUDE_DIR", "EIGEN_INCLUDE_DIR"]: + include_dir = os.environ.get(env_var) + if include_dir: + include_candidates.append(include_dir) + + for env_var in ["EIGEN3_PREFIX", "EIGEN_PREFIX"]: + prefix = os.environ.get(env_var) + if prefix: + include_candidates += [ + prefix, + os.path.join(prefix, "include"), + os.path.join(prefix, "include", "eigen3"), + ] + + if not (conda_prefix is None): + if system_name in ["Darwin", "Linux"]: + conda_include_path = os.path.join(conda_prefix, "include") + else: + conda_include_path = os.path.join(conda_prefix, "Library", "include") + include_candidates += [ + conda_include_path, + os.path.join(conda_include_path, "eigen3"), + ] + + if system_name == "Darwin": + brew_eigen_prefix = maybe_run_cmd("brew --prefix eigen") + if brew_eigen_prefix: + include_candidates += [ + os.path.join(brew_eigen_prefix, "include"), + os.path.join(brew_eigen_prefix, "include", "eigen3"), + ] + + if system_name in ["Darwin", "Linux"]: + include_candidates += [ + "/opt/homebrew/include/eigen3", + "/usr/local/include/eigen3", + "/usr/include/eigen3", + ] + + eigen_include_dirs = [] + for include_dir in include_candidates: + if has_eigen_headers(include_dir): + append_unique_dirs(eigen_include_dirs, [include_dir]) + + if not eigen_include_dirs: + raise RuntimeError( + "Eigen headers are not detected. " + "Set EIGEN3_INCLUDE_DIR to the directory containing 'Eigen/Core', " + "activate a conda environment containing eigen, " + "or install Homebrew and run 'brew install eigen'." + ) + + return eigen_include_dirs + + ParallelCompile("NPY_NUM_BUILD_JOBS").install() @@ -71,17 +146,7 @@ def run_cmd(cmd): system_name = platform.system() -# add include and include/eigen3 -if not (conda_prefix is None): - if system_name in ["Darwin", "Linux"]: - conda_include_path = os.path.join(conda_prefix, "include") - else: - conda_include_path = os.path.join(conda_prefix, "Library", "include") - eigen_include_path = os.path.join(conda_include_path, "eigen3") - include_dirs += [ - conda_include_path, - eigen_include_path, - ] +include_dirs += resolve_eigen_include_dirs(conda_prefix, system_name) if system_name == "Darwin": # if user provides OpenMP install prefix (containing include/ and lib/) @@ -188,4 +253,4 @@ def run_cmd(cmd): }, ext_modules=ext_modules, zip_safe=False, -) \ No newline at end of file +) diff --git a/tests/test_io.py b/tests/test_io.py index f06ea265..a1fc0fd4 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -116,4 +116,175 @@ def create_dense(calldata, ancestries, A, hap=None): my_dense = handler.to_dense() assert np.allclose(my_dense, dense) - os.remove(filename) \ No newline at end of file + os.remove(filename) + + +@pytest.mark.parametrize("read_mode", ["file", "mmap"]) +@pytest.mark.parametrize("n, s, A", [ + [1, 1, 1], + [200, 32, 4], + [2000, 3000, 7], + [1421, 927, 8], +]) +def test_io_snp_combine_r(n, s, A, read_mode, seed=0): + data = ad.data.snp_combine_r(n, s, A, seed=seed) + calldata = data["X"] + ancestries = data["ancestries"] + + # build expected dense: for each SNP j, + # - first column is the genotype calldata[:, j] + # - next A columns are the unphased dosage per ancestry + X_expected = np.zeros((n, s * (1 + A)), dtype=np.int8) + for j in range(s): + # genotype column + X_expected[:, j * (1 + A)] = calldata[:, j] + # ancestry dosage columns + for k in range(A): + dosage = ( + (ancestries[:, 2 * j] == k).astype(np.int8) + + (ancestries[:, 2 * j + 1] == k).astype(np.int8) + ) + X_expected[:, j * (1 + A) + 1 + k] = dosage + + filename = "/tmp/dummy_snp_combine_r.snpdat" + handler = ad.io.snp_combine_r(filename, read_mode=read_mode) + + # print matrix and dtype + print(f"calldata:\n{calldata}") + print(f"ancestries:\n{ancestries}") + print(f"calldata.dtype:\n{calldata.dtype}") + print(f"ancestries.dtype:\n{ancestries.dtype}") + + # write & read twice + w_bytes, _ = handler.write(calldata, ancestries, A, n_threads=2) + r_bytes1 = handler.read() + r_bytes2 = handler.read() + assert w_bytes == r_bytes1 == r_bytes2 + + # basic metadata + assert handler.rows == n + assert handler.snps == s + assert handler.ancestries == A + assert handler.cols == s * (1 + A) + + # full dense round‐trip + dense = handler.to_dense() + print(f"dense:\n{dense}") + print(f"X_expected:\n{X_expected}") + assert np.allclose(dense, X_expected) + + # nnz per column (non‐zero counts) + expected_nnz = np.sum(X_expected != 0, axis=0) + print(f"expected_nnz:\n{expected_nnz}") + print(f"handler.nnz:\n{handler.nnz}") + assert np.allclose(handler.nnz, expected_nnz) + + os.remove(filename) + + +@pytest.mark.parametrize("read_mode", ["file", "mmap"]) +@pytest.mark.parametrize("n, s, A", [ + [1, 1, 1], + [37, 13, 3], + [200, 32, 4], +]) +def test_io_snp_combine_s(n, s, A, read_mode, seed=0): + # Generate phased haplotype calldata (0/1 per hap) and phased ancestries + data = ad.data.snp_phased_ancestry(n, s, A, seed=seed) + calldata = data["X"] # shape (n, 2*s), entries in {0,1} + ancestries = data["ancestries"] # shape (n, 2*s), entries in [0,A) + + # Build expected dense (n, s*(2*A)) + X_expected = np.zeros((n, s * (2 * A)), dtype=np.int8) + for j in range(s): + c0 = calldata[:, 2*j] + c1 = calldata[:, 2*j + 1] + a0 = ancestries[:, 2*j] + a1 = ancestries[:, 2*j + 1] + for a in range(A): + # mutated hap counts per ancestry (first A columns) + mut = (c0 == 1).astype(np.int8) * (a0 == a).astype(np.int8) \ + + (c1 == 1).astype(np.int8) * (a1 == a).astype(np.int8) + X_expected[:, j * (2*A) + a] = mut + # ancestry dosage per ancestry (last A columns) + dos = (a0 == a).astype(np.int8) + (a1 == a).astype(np.int8) + X_expected[:, j * (2*A) + A + a] = dos + + filename = "/tmp/dummy_snp_combine_s.snpdat" + handler = ad.io.snp_combine_s(filename, read_mode=read_mode) + + # write & read twice + w_bytes, _ = handler.write(calldata, ancestries, A, n_threads=2) + r_bytes1 = handler.read() + r_bytes2 = handler.read() + assert w_bytes == r_bytes1 == r_bytes2 + + # basic metadata + assert handler.rows == n + assert handler.snps == s + assert handler.ancestries == A + assert handler.cols == s * (2 * A) + + # dense round-trip + dense = handler.to_dense(n_threads=2) + assert np.allclose(dense, X_expected) + + # nnz per column (non-zero sample counts) + expected_nnz = np.sum(X_expected != 0, axis=0) + assert np.allclose(handler.nnz, expected_nnz) + + os.remove(filename) + + +@pytest.mark.parametrize("read_mode", ["file", "mmap"]) +@pytest.mark.parametrize("n, s, A", [ + [37, 13, 3], + [200, 32, 4], +]) +def test_io_snp_combine_s_selected(n, s, A, read_mode, seed=1): + data = ad.data.snp_phased_ancestry(n, s, A, seed=seed) + calldata = data["X"] + ancestries = data["ancestries"] + + # Select a subset of ancestries (preserve order, unique) + if A >= 3: + selected = np.array([0, A-1], dtype=np.uint32) + else: + selected = np.array([0], dtype=np.uint32) + B = len(selected) + + # Build expected dense (n, s*(2*B)) restricted to selected ancestries + X_expected = np.zeros((n, s * (2 * B)), dtype=np.int8) + for j in range(s): + c0 = calldata[:, 2*j] + c1 = calldata[:, 2*j + 1] + a0 = ancestries[:, 2*j] + a1 = ancestries[:, 2*j + 1] + for bi, a in enumerate(selected): + # mutated block + mut = (c0 == 1).astype(np.int8) * (a0 == a).astype(np.int8) \ + + (c1 == 1).astype(np.int8) * (a1 == a).astype(np.int8) + X_expected[:, j * (2*B) + bi] = mut + # dosage block + dos = (a0 == a).astype(np.int8) + (a1 == a).astype(np.int8) + X_expected[:, j * (2*B) + B + bi] = dos + + filename = "/tmp/dummy_snp_combine_s_sel.snpdat" + handler = ad.io.snp_combine_s(filename, read_mode=read_mode) + w_bytes, _ = handler.write(calldata, ancestries, A, n_threads=2, selected_ancestries=selected) + r_bytes = handler.read() + assert w_bytes == r_bytes + + # Metadata reflects selection + assert handler.rows == n + assert handler.snps == s + assert handler.ancestries == B + assert handler.cols == s * (2 * B) + + dense = handler.to_dense(n_threads=2) + assert np.allclose(dense, X_expected) + + expected_nnz = np.sum(X_expected != 0, axis=0) + assert np.allclose(handler.nnz, expected_nnz) + + os.remove(filename) diff --git a/tests/test_matrix.py b/tests/test_matrix.py index 49ade58c..4cce196f 100644 --- a/tests/test_matrix.py +++ b/tests/test_matrix.py @@ -970,5 +970,149 @@ def test_naive_rsubset(n, p, subset_prop, dtype, seed=0): assert np.allclose(var, expected, atol=1e-6) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("read_mode", ["file", "mmap"]) +@pytest.mark.parametrize("n, s, A", [ + [10, 20, 4], + [1, 13, 3], + [144, 1, 2], + [10000, 1, 2], + [5, 130, 3], +]) +def test_naive_snp_combine_r(n, s, A, read_mode, dtype, seed=0): + """Round-trip IO ⟷ matrix tests for the SNP *unphased*, ancestry format. + + The test mirrors the existing SNP phased/unphased tests: + + 1. Generate a synthetic data-set using :pyfunc:`adelie.data.snp_combine_r`. + 2. Serialize it to disk via :pyclass:`adelie.io.snp_combine_r`. + 3. Wrap the on-disk representation as a matrix using the + :pyfunc:`adelie.matrix.snp_combine_r` helper. + 4. Verify that all naive-matrix primitives agree with a dense NumPy + reference via the shared ``run_naive`` routine. + """ + + def create_dense(calldata, ancestries, A): + n, s = calldata.shape + dense = np.zeros((n, s * (1 + A)), dtype=np.int8) + for j in range(s): + # genotype column + dense[:, j * (1 + A)] = calldata[:, j] + # ancestry dosage columns + for k in range(A): + dosage = ( + (ancestries[:, 2 * j] == k).astype(np.int8) + + (ancestries[:, 2 * j + 1] == k).astype(np.int8) + ) + dense[:, j * (1 + A) + 1 + k] = dosage + return dense + + # Reduce the *min_bytes* guard so that small test cases exercise all code + # paths irrespective of matrix size. + min_bytes = ad.configs.Configs.min_bytes + ad.configs.set_configs("min_bytes", 0) + + np.random.seed(seed) + data = ad.data.snp_combine_r(n, s, A, seed=seed) + + # Reference dense representation (float cast for arithmetic accuracy). + X = create_dense(data["X"], data["ancestries"], A).astype(dtype) + + filename = "/tmp/test_snp_combine_r.snpdat" + handler = ad.io.snp_combine_r(filename, read_mode) + handler.write(data["X"], data["ancestries"], A) + + # Construct the matrix wrapper. + cX = mod.snp_combine_r( + io=handler, + dtype=dtype, + n_threads=2, + ) + + # We can remove the temporary file immediately – the mmap handler keeps a + # reference if needed. + os.remove(filename) + + # Core comparison suite. + run_naive(X, cX, dtype) + + # Additional sanity-checks for *mean* / *var* helpers -------------------- + w = np.random.uniform(0, 1, cX.shape[0]).astype(dtype) + + mean = np.empty(cX.shape[1], dtype=dtype) + cX.mean(w, mean) + assert np.allclose(mean, 0) + + var = np.empty(cX.shape[1], dtype=dtype) + cX.var(mean, w, var) + assert np.allclose(var, 1) + + # Restore global configuration. + ad.configs.set_configs("min_bytes", min_bytes) + + # Reset to default settings -ad.configs.set_configs("min_bytes", None) \ No newline at end of file +ad.configs.set_configs("min_bytes", None) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("read_mode", ["file", "mmap"]) +@pytest.mark.parametrize("n, s, A", [ + [10, 20, 4], + [1, 7, 3], +]) +def test_naive_snp_combine_s(n, s, A, read_mode, dtype, seed=0): + """Round-trip IO ⟷ matrix tests for the SNP both-ancestry format.""" + + def create_dense(calldata, ancestries, A): + n, ss2 = calldata.shape + assert ss2 == 2 * s + dense = np.zeros((n, s * (2 * A)), dtype=np.int8) + for j in range(s): + c0 = calldata[:, 2*j] + c1 = calldata[:, 2*j + 1] + a0 = ancestries[:, 2*j] + a1 = ancestries[:, 2*j + 1] + for a in range(A): + # mutated block + mut = (c0 == 1).astype(np.int8) * (a0 == a).astype(np.int8) \ + + (c1 == 1).astype(np.int8) * (a1 == a).astype(np.int8) + dense[:, j*(2*A) + a] = mut + # dosage block + dos = (a0 == a).astype(np.int8) + (a1 == a).astype(np.int8) + dense[:, j*(2*A) + A + a] = dos + return dense + + # Reduce min_bytes guard so small tests exercise parallel code paths + min_bytes = ad.configs.Configs.min_bytes + ad.configs.set_configs("min_bytes", 0) + + np.random.seed(seed) + data = ad.data.snp_phased_ancestry(n, s, A, seed=seed) + Xdense = create_dense(data["X"], data["ancestries"], A).astype(dtype) + + filename = "/tmp/test_snp_combine_s.snpdat" + handler = ad.io.snp_combine_s(filename, read_mode) + handler.write(data["X"], data["ancestries"], A) + + cX = mod.snp_combine_s( + io=handler, + dtype=dtype, + n_threads=2, + ) + + os.remove(filename) + + run_naive(Xdense, cX, dtype) + + # mean/var follow same semantics as unphased ancestry helpers + w = np.random.uniform(0, 1, cX.shape[0]).astype(dtype) + mean = np.empty(cX.shape[1], dtype=dtype) + cX.mean(w, mean) + assert np.allclose(mean, 0) + + var = np.empty(cX.shape[1], dtype=dtype) + cX.var(mean, w, var) + assert np.allclose(var, 1) + + ad.configs.set_configs("min_bytes", min_bytes) \ No newline at end of file diff --git a/tests/test_solver.py b/tests/test_solver.py index b3a0c292..bc7da9e9 100644 --- a/tests/test_solver.py +++ b/tests/test_solver.py @@ -818,6 +818,90 @@ def test_solve_gaussian_snp_unphased( os.remove(filename) +@pytest.mark.filterwarnings("ignore: Detected matrix to be C-contiguous.") +@pytest.mark.parametrize("n, p", [ + [10, 4], + [10, 100], + [100, 23], + [100, 100], + [100, 10000], +]) +def test_solve_gaussian_snp_combine_r( + n, p, + A=8, intercept=True, alpha=1, sparsity=0.5, seed=0, n_threads=2, +): + """Compares Gaussian solver outputs for SNP-unphased-ancestry matrices + against their dense equivalents. + + The test mirrors ``test_solve_gaussian_snp_phased_ancestry`` but for the + unphased ancestry format (one ancestry-dosage column per ancestry). + """ + # ------------------------------------------------------------------ + # 1. Generate synthetic dataset and serialise it to disk + # ------------------------------------------------------------------ + test_data = ad.data.snp_combine_r(n=n, s=p, A=A, sparsity=sparsity, seed=seed) + filename = "/tmp/test_snp_combine_r.snpdat" + handler = ad.io.snp_combine_r(filename) + handler.write(test_data["X"], test_data["ancestries"], A, n_threads) + Xs = [ + ad.matrix.snp_combine_r( + io=handler, + dtype=np.float64, + n_threads=n_threads, + ) + ] + handler.read() + + X, y = handler.to_dense(n_threads), test_data.pop("glm").y + + weights = np.random.uniform(1, 2, n) + weights /= np.sum(weights) + + test_data["constraints"] = None + test_data["y"] = y + test_data["weights"] = weights + test_data["offsets"] = np.zeros(n) + test_data["alpha"] = alpha + test_data["X_means"] = np.sum(weights[:, None] * X, axis=0) + test_data["y_mean"] = np.sum(weights * y) + X_c = X - intercept * test_data["X_means"][None] + y_c = y - test_data["y_mean"] * intercept + test_data["y_var"] = np.sum(weights * y_c ** 2) + test_data["resid"] = y_c + test_data["resid_sum"] = np.sum(weights * test_data["resid"]) + test_data["screen_set"] = np.arange(p)[(test_data["penalty"] <= 0) | (alpha <= 0)] + test_data["screen_beta"] = np.zeros(np.sum(test_data["group_sizes"][test_data["screen_set"]])) + test_data["screen_is_active"] = np.zeros(test_data["screen_set"].shape[0], dtype=bool) + test_data["active_set_size"] = 0 + test_data["active_set"] = np.empty(p, dtype=int) + test_data["grad"] = X_c.T @ (weights * test_data["resid"]) + test_data["rsq"] = 0 + test_data["lmda"] = np.inf + test_data["tol"] = 1e-7 + test_data["n_threads"] = n_threads + + test_data.pop("ancestries") + + for Xpy in Xs: + test_data["X"] = Xpy + state_special = ad.state.gaussian_naive(**test_data).solve() + test_data["X"] = ad.matrix.dense( + X.astype(np.float64), + method="naive", + n_threads=n_threads + ) + state_dense = ad.state.gaussian_naive(**test_data).solve() + + assert np.allclose(state_special.lmdas, state_dense.lmdas) + assert np.allclose(state_special.devs, state_dense.devs) + assert np.allclose(state_special.intercepts, state_dense.intercepts) + assert np.allclose( + state_special.betas.toarray(), state_dense.betas.toarray() + ) + + os.remove(filename) + + @pytest.mark.filterwarnings("ignore: Detected matrix to be C-contiguous.") @pytest.mark.parametrize("n, p", [ [10, 4],