{"id":13520972,"url":"https://github.com/mancusolab/traceax","last_synced_at":"2025-12-14T08:30:40.096Z","repository":{"id":228900785,"uuid":"775179907","full_name":"mancusolab/traceax","owner":"mancusolab","description":"Stochastic trace estimation using JAX","archived":false,"fork":false,"pushed_at":"2025-08-20T09:25:45.000Z","size":739,"stargazers_count":15,"open_issues_count":3,"forks_count":3,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-08-20T11:37:10.315Z","etag":null,"topics":["jax","machine-learn","python3","statistics","trace-est"],"latest_commit_sha":null,"homepage":"https://mancusolab.github.io/traceax","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mancusolab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2024-03-20T22:50:10.000Z","updated_at":"2025-08-20T09:25:48.000Z","dependencies_parsed_at":"2024-04-29T18:46:55.733Z","dependency_job_id":"25447e1d-5c8b-4729-b18d-98f313147667","html_url":"https://github.com/mancusolab/traceax","commit_stats":null,"previous_names":["mancusolab/traceax"],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/mancusolab/traceax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mancusolab%2Ftraceax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mancusolab%2Ftraceax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mancusolab%2Ftraceax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mancusolab%2Ftraceax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mancusolab","download_url":"https://codeload.github.com/mancusolab/traceax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mancusolab%2Ftraceax/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":27723235,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-12-14T02:00:11.348Z","response_time":56,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["jax","machine-learn","python3","statistics","trace-est"],"created_at":"2024-08-01T06:00:25.330Z","updated_at":"2025-12-14T08:30:40.085Z","avatar_url":"https://github.com/mancusolab.png","language":"Python","funding_links":[],"categories":["Models and Projects","Libraries"],"sub_categories":["Up and Coming Libraries"],"readme":"[![Documentation-webpage](https://img.shields.io/badge/Docs-Available-brightgreen)](https://mancusolab.github.io/traceax/)\n[![PyPI-Server](https://img.shields.io/pypi/v/traceax.svg)](https://pypi.org/project/traceax/)\n[![Github](https://img.shields.io/github/stars/mancusolab/traceax?style=social)](https://github.com/mancusolab/traceax)\n[![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)\n[![Project generated with Hatch](https://img.shields.io/badge/%F0%9F%A5%9A-Hatch-4051b5.svg)](https://github.com/pypa/hatch)\n\n# traceax\n``traceax`` is a Python library to perform stochastic trace estimation for linear operators. Namely,\ngiven a square linear operator \u003ci\u003e`A`\u003c/i\u003e, ``traceax`` provides flexible routines that estimate,\n\n$$\\text{trace}(\\mathbf{A}) = \\sum_i \\mathbf{A}_{ii},$$\n\nusing only matrix-vector products. ``traceax`` is heavily inspired by\n[lineax](https://github.com/patrick-kidger/lineax) as well as\n[XTrace](https://github.com/eepperly/XTrace).\n\n  [**Installation**](#installation)\n  | [**Example**](#get-started-with-example)\n  | [**Documentation**](#documentation)\n  | [**Notes**](#notes)\n  | [**Support**](#support)\n  | [**Other Software**](#other-software)\n\n------------------\n\n\u003e [!NOTE]\n\u003e For preprint, please see: \u003cbr/\u003e\n\u003e Nahid, A.A., Serafin, L., Mancuso, N. (2025). \u003ci\u003etraceax\u003c/i\u003e: a JAX-based framework for stochastic trace estimation. bioRxiv \n\u003e [doi.org/10.1101/2025.07.14.662216](https://doi.org/10.1101/2025.07.14.662216)\n\n\u003e [!IMPORTANT]\n\u003e To reproduce the analyses in the preprint: \u003cbr/\u003e\n\u003e [github.com/mancusolab/traceax_analysis](https://github.com/mancusolab/traceax_analysis)\n\n------------------\n\n## Installation\n\nUsers can directly install from `pip`:\n\n``` bash\npip install traceax\n```\n\nOr, users can download the latest repository and then use `pip`:\n\n```\ngit clone https://github.com/mancusolab/traceax.git\ncd traceax\npip install .\n```\n\n## Get Started with Example\n\n```python\nimport jax.numpy as jnp\nimport jax.random as rdm\nimport lineax as lx\n\nimport traceax as tx\n\n# simulate simple symmetric matrix with exponential eigenvalue decay\nseed = 0\nN = 1000\nkey = rdm.PRNGKey(seed)\nkey, xkey = rdm.split(key)\n\nX = rdm.normal(xkey, (N, N))\nQ, R = jnp.linalg.qr(X)\nU = jnp.power(0.7, jnp.arange(N))\nA = (Q * U) @ Q.T\n\n# should be numerically close\nprint(jnp.trace(A))  # 3.3333323\nprint(jnp.sum(U))  # 3.3333335\n\n# setup linear operator\noperator = lx.MatrixLinearOperator(A)\n\n# number of matrix vector operators\nk = 25\n\n# split key for estimators\nkey, key1, key2, key3, key4 = rdm.split(key, 5)\n\n# Hutchinson estimator; default samples Rademacher {-1,+1}\nhutch = tx.HutchinsonEstimator()\nprint(hutch.estimate(key1, operator, k))  # (Array(3.4099615, dtype=float32), {})\n\n# Hutch++ estimator; default samples Rademacher {-1,+1}\nhpp = tx.HutchPlusPlusEstimator()\nprint(hpp.estimate(key2, operator, k))  # (Array(3.3033807, dtype=float32), {})\n\n# XTrace estimator; default samples uniformly on n-Sphere\nxt = tx.XTraceEstimator()\nprint(xt.estimate(key3, operator, k))  # (Array(3.3271673, dtype=float32), {'std.err': Array(0.01717775, dtype=float32)})\n\n# XNysTrace estimator; Improved performance for NSD/PSD trace estimates\noperator = lx.TaggedLinearOperator(operator, lx.positive_semidefinite_tag)\nnt = tx.XNysTraceEstimator()\nprint(nt.estimate(key4, operator, k))  # (Array(3.3297246, dtype=float32), {'std.err': Array(0.00042093, dtype=float32)})\n```\n\n## Documentation\nDocumentation is available at [here](https://mancusolab.github.io/traceax/).\n\n## Notes\n\n-   `traceax` uses [JAX](https://github.com/google/jax) with [Just In\n    Time](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html)\n    compilation to achieve high-speed computation. However, there are\n    some [issues](https://github.com/google/jax/issues/5501) for JAX\n    with Mac M1 chip. To solve this, users need to initiate conda using\n    [miniforge](https://github.com/conda-forge/miniforge), and then\n    install `traceax` using `pip` in the desired environment.\n\n\n## Support\n\nPlease report any bugs or feature requests in the [Issue\nTracker](https://github.com/mancusolab/traceax/issues). If users have\nany questions or comments, please contact Abdullah Al Nahid (\u003calnahid@usc.edu\u003e) or\nNicholas Mancuso (\u003cnmancuso@usc.edu\u003e).\n\n## Other Software\n\nFeel free to use other software developed by [Mancuso\nLab](https://www.mancusolab.com/):\n\n-   [SuShiE](https://github.com/mancusolab/sushie): a Bayesian\n    fine-mapping framework for molecular QTL data across multiple\n    ancestries.\n-   [MA-FOCUS](https://github.com/mancusolab/ma-focus): a Bayesian\n    fine-mapping framework using\n    [TWAS](https://www.nature.com/articles/ng.3506) statistics across\n    multiple ancestries to identify the causal genes for complex traits.\n-   [SuSiE-PCA](https://github.com/mancusolab/susiepca): a scalable\n    Bayesian variable selection technique for sparse principal component\n    analysis\n-   [twas_sim](https://github.com/mancusolab/twas_sim): a Python\n    software to simulate [TWAS](https://www.nature.com/articles/ng.3506)\n    statistics.\n-   [FactorGo](https://github.com/mancusolab/factorgo): a scalable\n    variational factor analysis model that learns pleiotropic factors\n    from GWAS summary statistics.\n-   [HAMSTA](https://github.com/tszfungc/hamsta): a Python software to\n    estimate heritability explained by local ancestry data from\n    admixture mapping summary statistics.\n\n------------------------------------------------------------------------\n\n``traceax`` is distributed under the terms of the\n[Apache-2.0 license](https://spdx.org/licenses/Apache-2.0.html).\n\n\n------------------------------------------------------------------------\n\nThis project has been set up using Hatch. For details and usage\ninformation on Hatch see \u003chttps://github.com/pypa/hatch\u003e.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmancusolab%2Ftraceax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmancusolab%2Ftraceax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmancusolab%2Ftraceax/lists"}