fastmixture2 2.0.3__tar.gz → 2.0.4__tar.gz

{fastmixture2-2.0.3/fastmixture2.egg-info → fastmixture2-2.0.4}/PKG-INFO

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: fastmixture2
-Version: 2.0.3
-Summary: Fast and scalable memory-aware maximum-likelihood ancestry estimation
-Author-email: Jonas Meisner <meisnerucph@gmail.com>
+Version: 2.0.4
+Summary: Biobank-scale ancestry estimation in unrelated individuals
+Author-email: "J. Meisner" <meisnerucph@gmail.com>, "L. Thunegard" <lisathunegard@icloud.com>
 License-Expression: GPL-3.0-only
 Project-URL: Homepage, https://github.com/MeisnerLab/fastmixture2
 Project-URL: Documentation, https://github.com/MeisnerLab/fastmixture2#readme
@@ -32,13 +32,14 @@ Dynamic: license-file
 # fastmixture2
 
 [![Tests](https://github.com/MeisnerLab/fastmixture2/actions/workflows/ci.yml/badge.svg)](https://github.com/MeisnerLab/fastmixture2/actions/workflows/ci.yml)
-[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/fastmixture2.svg)](https://pypi.org/project/fastmixture2/)
+[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://github.com/MeisnerLab/fastmixture2/blob/main/LICENSE)
 
 `fastmixture2` estimates ancestry proportions (`Q`) and ancestral allele
 frequencies (`P`) by maximum likelihood. It reads SNP-major PLINK genotypes
 or ANGSD Beagle genotype likelihoods. Unsupervised initialization uses SVD/ALS
-for K > 1. All modes use safeguarded SQP/ZAL updates on CPU or, with `--gpu`,
-on NVIDIA CUDA.
+for K > 1, with subset likelihood refinement for large panels.
+All modes use safeguarded SQP/ZAL updates on CPU or, with `--gpu`, on NVIDIA CUDA.
 
 ## Installation
 
@@ -58,7 +59,8 @@ Source builds need a C compiler with OpenMP. On macOS, install OpenMP with
 `brew install libomp`. On HPC, build on the compute-node type used for your
 analyses.
 
-With Conda, the [CPU](environment-cpu.yml) and [GPU](environment-gpu.yml)
+With Conda, the [CPU](https://github.com/MeisnerLab/fastmixture2/blob/main/environment-cpu.yml)
+and [GPU](https://github.com/MeisnerLab/fastmixture2/blob/main/environment-gpu.yml)
 environments retain native CPU tuning. Run from the cloned repository:
 
 ```bash
@@ -69,7 +71,7 @@ conda activate fastmixture2-cpu
 For GPU use, select `environment-gpu.yml` and activate `fastmixture2-gpu`.
 
 PyPI wheels cover Linux (glibc) and macOS 14 or newer, on x86-64 and ARM64.
-They are portable and select SIMD kernels when available:
+They use AVX2 on supported x86-64 CPUs and NEON on ARM64, including Apple Silicon:
 
 ```bash
 python -m pip install fastmixture2
@@ -132,20 +134,39 @@ fastmixture2 --bfile cohort --K 5 --threads 8 --gpu --out cohort.gpu
 CPU and CUDA support `1 <= K <= 128` for both input formats in all modes.
 CUDA keeps input data on the GPU when they fit and streams in blocks otherwise.
 Small CPU/GPU differences can arise from floating-point arithmetic.
-At K > 1, CUDA SVD starts at 256 samples. Large ALS factors at K <= 8 use
-CUDA ALS with float64 contractions. Smaller or high-K factors retain CPU ALS.
-Resident PLINK CUDA fits use float32 SQP statistics while gains remain material,
-then switch to float64 refinement. There is no fixed float32 iteration cap.
+
+At K > 1, unsupervised initialization uses 10 SVD power iterations on all SNPs
+up to 50,000, or a genome-wide 50,000-SNP subset otherwise. Larger panels refine
+subset Q/P with SQP/ZAL, then fit full-marker P by least squares at fixed Q and
+retain the refined subset rows. This applies to both input formats, including
+missing genotypes.
+Panels up to 50,000 SNPs proceed directly from SVD/ALS to full-data refinement,
+without a duplicate subset fit or projection. For PLINK, subset refinement is
+skipped if any sample has no observed subset genotypes. Exact CV folds and
+`--iter 0` retain SVD/ALS without subset likelihood refinement.
+
+Subsets exceeding 512 MiB on CPU or 256 MiB on CUDA are built in bounded blocks
+into a temporary memory-mapped file. CUDA streams the subset when it does not
+fit in VRAM. Temporary files use `TMPDIR`, which can point to disk-backed scratch
+storage.
+
+CUDA SVD starts at 256 samples. Both devices use CPU ALS with float64
+contractions. Fixed-Q P projection uses float64 products on CPU and float32
+products on CUDA.
+
+Resident and streamed PLINK CUDA fits use float32 SQP statistics while gains
+remain material, then switch to float64 refinement. There is no fixed float32
+iteration cap. Tile products use float32, with float64 sums across tiles.
 Parameters, constrained solves and likelihood checks stay in float64.
-PLINK cuBLAS uses float32 tile products with float64 accumulation.
 
 Genotype-likelihood fits use Q-first SQP/ZAL on both CPU and CUDA, with exact
-EM/QN rescue and an independent EM check before convergence. Resident CUDA
-fits start with float32 probabilities, derivatives and statistics,
+EM/QN rescue and an independent EM check before convergence. CUDA fits start
+with float32 probabilities, derivatives and statistics in either storage mode,
 then switch to float64 refinement. Near frequency bounds, GL probabilities
 and derivatives use float64. cuBLAS tile sums, parameters, solves and scores stay
 in float64.
-CPU and CUDA start directly with SQP. `--iter` limits refinement.
+CPU and CUDA start directly with SQP and use ZAL history depth 3.
+`--iter` and `--tole` apply to each SQP/ZAL stage.
 EM skips missing genotypes and exactly uniform likelihood triplets, with
 normalization over informative entries. Nonuniform likelihoods are retained.
 Full likelihood scores retain the constant contribution from uniform triplets.
@@ -157,13 +178,16 @@ use 4096-sample shards with a separate 4 GiB cap. Optional update caches
 remain capped at 64 MiB. Input storage, factors and optimizer history are
 additional allocations.
 
-Resident CUDA uses native kernels at K = 1..3 and tiled cuBLAS for Q and P
-at K = 4..16 during float32 refinement, regardless of input format or dataset size.
-Tiles share at most 64 MiB of scratch and keep float32 copies of P and Q.
-Solves and likelihood checks remain in float64. At K <= 16, final refinement
-and streamed fits use native float64 kernels, except resident GL P updates at
-K = 8..16 with up to 1024 samples. Those share the tiled cuBLAS P routine and
-keep combined scratch within 64 MiB.
+At K <= 16, CUDA fits use native kernels at K = 1..3
+and tiled cuBLAS at K = 4..16 during float32 refinement. Resident tiles share
+64 MiB of scratch and keep float32 copies of P and Q. At K = 8..16,
+resident mixed fits reuse frequency scratch for cooperative float64 solves.
+Tiles that cannot hold the solve tables use the native solve. Streamed fits keep
+float32 copies within tiles and reuse their float64 Q-statistics buffers, with
+up to 128 MiB of compute scratch.
+Final refinement uses native float64 kernels, except
+resident GL P updates at K = 8..16 with up to 1024 samples. Those share the tiled
+cuBLAS P routine and keep combined scratch within 64 MiB.
 Unsupervised resident PLINK fits at K = 10..16 keep the first eight iterations
 fresh, then alternate Q-Hessian reuse during float32 refinement. The cache is
 capped at 64 MiB. Weak cached proposals retry with fresh curvature, and
@@ -172,8 +196,8 @@ float64 refinement always uses fresh curvature.
 At K > 16, CUDA uses tiled cuBLAS statistics and a cooperative active-set
 solver for both input formats. Resident and bootstrap maps share 64 MiB of
 scratch. Streamed maps allow up to 128 MiB of compute scratch and 512 MiB of
-Q statistics, with free-VRAM guards. Curvature stays fresh. Larger K still
-increases compute cost, especially for the quadratic solves.
+Q statistics, with free-VRAM guards. Streamed tiles adapt to the panel width,
+with at most 8192 rows within the same memory caps. Curvature stays fresh.
 
 The defaults are `--threads 1`, `--seed 42`, `--iter 1000` and `--tole 1e-10`.
 `--threads` caps CPU computation, including BLAS and GPU host copies.
@@ -333,4 +357,4 @@ fastmixture2 does not create filesystem locks or require private lock directorie
 ## Authors and license
 
 Lisa Thunegard and Jonas Meisner. Developed at **MeisnerLab** and distributed
-under the [GNU General Public License v3.0](LICENSE).
+under the [GNU General Public License v3.0](https://github.com/MeisnerLab/fastmixture2/blob/main/LICENSE).

{fastmixture2-2.0.3 → fastmixture2-2.0.4}/README.md

@@ -1,13 +1,14 @@
 # fastmixture2
 
 [![Tests](https://github.com/MeisnerLab/fastmixture2/actions/workflows/ci.yml/badge.svg)](https://github.com/MeisnerLab/fastmixture2/actions/workflows/ci.yml)
-[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](LICENSE)
+[![PyPI](https://img.shields.io/pypi/v/fastmixture2.svg)](https://pypi.org/project/fastmixture2/)
+[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://github.com/MeisnerLab/fastmixture2/blob/main/LICENSE)
 
 `fastmixture2` estimates ancestry proportions (`Q`) and ancestral allele
 frequencies (`P`) by maximum likelihood. It reads SNP-major PLINK genotypes
 or ANGSD Beagle genotype likelihoods. Unsupervised initialization uses SVD/ALS
-for K > 1. All modes use safeguarded SQP/ZAL updates on CPU or, with `--gpu`,
-on NVIDIA CUDA.
+for K > 1, with subset likelihood refinement for large panels.
+All modes use safeguarded SQP/ZAL updates on CPU or, with `--gpu`, on NVIDIA CUDA.
 
 ## Installation
 
@@ -27,7 +28,8 @@ Source builds need a C compiler with OpenMP. On macOS, install OpenMP with
 `brew install libomp`. On HPC, build on the compute-node type used for your
 analyses.
 
-With Conda, the [CPU](environment-cpu.yml) and [GPU](environment-gpu.yml)
+With Conda, the [CPU](https://github.com/MeisnerLab/fastmixture2/blob/main/environment-cpu.yml)
+and [GPU](https://github.com/MeisnerLab/fastmixture2/blob/main/environment-gpu.yml)
 environments retain native CPU tuning. Run from the cloned repository:
 
 ```bash
@@ -38,7 +40,7 @@ conda activate fastmixture2-cpu
 For GPU use, select `environment-gpu.yml` and activate `fastmixture2-gpu`.
 
 PyPI wheels cover Linux (glibc) and macOS 14 or newer, on x86-64 and ARM64.
-They are portable and select SIMD kernels when available:
+They use AVX2 on supported x86-64 CPUs and NEON on ARM64, including Apple Silicon:
 
 ```bash
 python -m pip install fastmixture2
@@ -101,20 +103,39 @@ fastmixture2 --bfile cohort --K 5 --threads 8 --gpu --out cohort.gpu
 CPU and CUDA support `1 <= K <= 128` for both input formats in all modes.
 CUDA keeps input data on the GPU when they fit and streams in blocks otherwise.
 Small CPU/GPU differences can arise from floating-point arithmetic.
-At K > 1, CUDA SVD starts at 256 samples. Large ALS factors at K <= 8 use
-CUDA ALS with float64 contractions. Smaller or high-K factors retain CPU ALS.
-Resident PLINK CUDA fits use float32 SQP statistics while gains remain material,
-then switch to float64 refinement. There is no fixed float32 iteration cap.
+
+At K > 1, unsupervised initialization uses 10 SVD power iterations on all SNPs
+up to 50,000, or a genome-wide 50,000-SNP subset otherwise. Larger panels refine
+subset Q/P with SQP/ZAL, then fit full-marker P by least squares at fixed Q and
+retain the refined subset rows. This applies to both input formats, including
+missing genotypes.
+Panels up to 50,000 SNPs proceed directly from SVD/ALS to full-data refinement,
+without a duplicate subset fit or projection. For PLINK, subset refinement is
+skipped if any sample has no observed subset genotypes. Exact CV folds and
+`--iter 0` retain SVD/ALS without subset likelihood refinement.
+
+Subsets exceeding 512 MiB on CPU or 256 MiB on CUDA are built in bounded blocks
+into a temporary memory-mapped file. CUDA streams the subset when it does not
+fit in VRAM. Temporary files use `TMPDIR`, which can point to disk-backed scratch
+storage.
+
+CUDA SVD starts at 256 samples. Both devices use CPU ALS with float64
+contractions. Fixed-Q P projection uses float64 products on CPU and float32
+products on CUDA.
+
+Resident and streamed PLINK CUDA fits use float32 SQP statistics while gains
+remain material, then switch to float64 refinement. There is no fixed float32
+iteration cap. Tile products use float32, with float64 sums across tiles.
 Parameters, constrained solves and likelihood checks stay in float64.
-PLINK cuBLAS uses float32 tile products with float64 accumulation.
 
 Genotype-likelihood fits use Q-first SQP/ZAL on both CPU and CUDA, with exact
-EM/QN rescue and an independent EM check before convergence. Resident CUDA
-fits start with float32 probabilities, derivatives and statistics,
+EM/QN rescue and an independent EM check before convergence. CUDA fits start
+with float32 probabilities, derivatives and statistics in either storage mode,
 then switch to float64 refinement. Near frequency bounds, GL probabilities
 and derivatives use float64. cuBLAS tile sums, parameters, solves and scores stay
 in float64.
-CPU and CUDA start directly with SQP. `--iter` limits refinement.
+CPU and CUDA start directly with SQP and use ZAL history depth 3.
+`--iter` and `--tole` apply to each SQP/ZAL stage.
 EM skips missing genotypes and exactly uniform likelihood triplets, with
 normalization over informative entries. Nonuniform likelihoods are retained.
 Full likelihood scores retain the constant contribution from uniform triplets.
@@ -126,13 +147,16 @@ use 4096-sample shards with a separate 4 GiB cap. Optional update caches
 remain capped at 64 MiB. Input storage, factors and optimizer history are
 additional allocations.
 
-Resident CUDA uses native kernels at K = 1..3 and tiled cuBLAS for Q and P
-at K = 4..16 during float32 refinement, regardless of input format or dataset size.
-Tiles share at most 64 MiB of scratch and keep float32 copies of P and Q.
-Solves and likelihood checks remain in float64. At K <= 16, final refinement
-and streamed fits use native float64 kernels, except resident GL P updates at
-K = 8..16 with up to 1024 samples. Those share the tiled cuBLAS P routine and
-keep combined scratch within 64 MiB.
+At K <= 16, CUDA fits use native kernels at K = 1..3
+and tiled cuBLAS at K = 4..16 during float32 refinement. Resident tiles share
+64 MiB of scratch and keep float32 copies of P and Q. At K = 8..16,
+resident mixed fits reuse frequency scratch for cooperative float64 solves.
+Tiles that cannot hold the solve tables use the native solve. Streamed fits keep
+float32 copies within tiles and reuse their float64 Q-statistics buffers, with
+up to 128 MiB of compute scratch.
+Final refinement uses native float64 kernels, except
+resident GL P updates at K = 8..16 with up to 1024 samples. Those share the tiled
+cuBLAS P routine and keep combined scratch within 64 MiB.
 Unsupervised resident PLINK fits at K = 10..16 keep the first eight iterations
 fresh, then alternate Q-Hessian reuse during float32 refinement. The cache is
 capped at 64 MiB. Weak cached proposals retry with fresh curvature, and
@@ -141,8 +165,8 @@ float64 refinement always uses fresh curvature.
 At K > 16, CUDA uses tiled cuBLAS statistics and a cooperative active-set
 solver for both input formats. Resident and bootstrap maps share 64 MiB of
 scratch. Streamed maps allow up to 128 MiB of compute scratch and 512 MiB of
-Q statistics, with free-VRAM guards. Curvature stays fresh. Larger K still
-increases compute cost, especially for the quadratic solves.
+Q statistics, with free-VRAM guards. Streamed tiles adapt to the panel width,
+with at most 8192 rows within the same memory caps. Curvature stays fresh.
 
 The defaults are `--threads 1`, `--seed 42`, `--iter 1000` and `--tole 1e-10`.
 `--threads` caps CPU computation, including BLAS and GPU host copies.
@@ -302,4 +326,4 @@ fastmixture2 does not create filesystem locks or require private lock directorie
 ## Authors and license
 
 Lisa Thunegard and Jonas Meisner. Developed at **MeisnerLab** and distributed
-under the [GNU General Public License v3.0](LICENSE).
+under the [GNU General Public License v3.0](https://github.com/MeisnerLab/fastmixture2/blob/main/LICENSE).

fastmixture2-2.0.4/fastmixture2/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.0.4"

{fastmixture2-2.0.3 → fastmixture2-2.0.4}/fastmixture2/defaults.py

@@ -14,9 +14,9 @@ LIKELIHOOD_CACHE_HEADER_BYTES = 256
 LIKELIHOOD_CACHE_VERSION = 5
 LIKELIHOOD_FREQUENCY_VERSION = 2
 
-# SVD/ALS initialization
-INITIALIZATION_MARKERS = 100_000
-SVD_POWER_ITERATIONS = 5
+# Initialization
+INITIALIZATION_MARKERS = 50_000
+SVD_POWER_ITERATIONS = 10
 SVD_CHUNK_MARKERS = 4_096
 SVD_WORKSPACE_BYTES = 512 << 20
 GPU_SVD_WORKSPACE_BYTES = 256 << 20

{fastmixture2-2.0.3 → fastmixture2-2.0.4}/fastmixture2/em.pyx

@@ -1244,15 +1244,17 @@ cpdef int accelPacked(
     return 0
 
 
-### Form one weighted packed marker-bootstrap EM map
-cdef int _batchPacked(
+### Form one multiplicity-weighted packed marker-bootstrap EM map
+cpdef int batchPackedW(
         const u8[:, ::1] B, const f64[:, ::1] P,
         f64[:, ::1] P1, const f64[:, ::1] Q,
-        f64[:, ::1] T, const u32[::1] s, const u32* w,
-        const Py_ssize_t N, const bint masked, const Py_ssize_t fold,
-        const Py_ssize_t folds, const uint64_t seed
+        f64[:, ::1] T, const u32[::1] s, const u32[::1] w,
+        const Py_ssize_t N, Py_ssize_t fold=-1,
+        Py_ssize_t folds=1, uint64_t seed=0
     ) except -1 nogil:
     cdef:
+        bint masked = fold >= 0
+        const u32* wt = &w[0]
         u8* g_thr
         u8* g_all
         const f64* p
@@ -1290,7 +1292,7 @@ cdef int _batchPacked(
         for j in prange(M, schedule="static"):
             marker = s[j]
             p = &P[marker, 0]
-            if w[j] == 0:
+            if wt[j] == 0:
                 for k in range(K):
                     P1[marker, k] = p[k]
                 continue
@@ -1302,7 +1304,7 @@ cdef int _batchPacked(
             else:
                 _unpackPackedRow(&B[marker, 0], g_thr, N)
             if not _specializedEmRowW(
-                    g_thr, p, &Q[0, 0], &P1[marker, 0], q_thr, w[j], N, K
+                    g_thr, p, &Q[0, 0], &P1[marker, 0], q_thr, wt[j], N, K
                 ):
                 for i in range(N):
                     if g_thr[i] == 9:
@@ -1311,7 +1313,7 @@ cdef int _batchPacked(
                     h = _computeH(p, q, K)
                     _innerJW(
                         p, q, &p_thr[0], &p_thr[K], &q_thr[i * K],
-                        g_thr[i], w[j], K, h
+                        g_thr[i], wt[j], K, h
                     )
                 _outerAccelP(p, &P1[marker, 0], &p_thr[0], &p_thr[K], K)
     for x in prange(L, schedule="static"):
@@ -1323,17 +1325,6 @@ cdef int _batchPacked(
     return 0
 
 
-### Form one multiplicity-weighted packed marker-bootstrap EM map
-cpdef int batchPackedW(
-        const u8[:, ::1] B, const f64[:, ::1] P, f64[:, ::1] P1,
-        const f64[:, ::1] Q, f64[:, ::1] T,
-        const u32[::1] s, const u32[::1] w,
-        const Py_ssize_t N, Py_ssize_t fold=-1,
-        Py_ssize_t folds=1, uint64_t seed=0
-    ) except -1 nogil:
-    return _batchPacked(B, P, P1, Q, T, s, &w[0], N, fold >= 0, fold, folds, seed)
-
-
 ### Form one packed fixed-P Q map
 cpdef int stepQPacked(
         const u8[:, ::1] B, const f64[:, ::1] P,

{fastmixture2-2.0.3 → fastmixture2-2.0.4}/fastmixture2/fastmixture2.py

@@ -805,6 +805,13 @@ def _run(args, stack):
     else:
         boot_h, b_chr = utils.bootBlock(args.bfile, M, args.bootstrap_block)
 
+    run = {
+        "iter": args.iter,
+        "tole": args.tole,
+        "check": check,
+        "gpu": use_gpu,
+        "log": runtime.append,
+    }
     print("\nInitialization:")
     resumed = pops = y = None
     if args.resume is not None:
@@ -861,6 +868,7 @@ def _run(args, stack):
             P1, Q = resumed[:2]
             if P1 is not None and not np.allclose(P1, P0, atol=1e-12, rtol=0.0):
                 raise ValueError("Checkpoint frequencies do not match --projection")
+            del P1
             P = P0
         del P0
         sec = time() - ts
@@ -893,10 +901,12 @@ def _run(args, stack):
             elif use_gpu:
                 from fastmixture2 import gpu
 
-                P, Q = gpu.initializeUnsupervised(G, args.K, rng)
+                P, Q = gpu.initializeUnsupervised(G, args.K, rng, run=run)
             else:
-                P, Q = utils.initializeUnsupervised(G, args.K, rng)
+                P, Q = utils.initializeUnsupervised(G, args.K, rng, run=run)
             sec = time() - ts
+            if "init" in run:
+                label = "SVD/ALS and subset refinement"
             functions.printTiming(f"{label}.", sec, replace=True)
             runtime.append(f"{label}: {sec:.1f} s")
         else:
@@ -909,14 +919,6 @@ def _run(args, stack):
     if inputKey(args) != src_key:
         raise ValueError("Analysis inputs changed while reading")
 
-    # Build the optimizer settings
-    run = {
-        "iter": args.iter,
-        "tole": args.tole,
-        "check": check,
-        "gpu": use_gpu,
-        "log": runtime.append,
-    }
     if resumed is not None:
         if not use_gpu:
             like = functions._loglike(G, q_nrm)
@@ -929,6 +931,7 @@ def _run(args, stack):
 
         # CUDA scores resumed estimates after their single host-to-device transfer
         run["off"] = resumed[3]
+    del resumed
     ckpt = args.checkpoint or args.resume
     if ckpt is not None:
         ### Persist the latest accepted optimizer state