nqxpack 0.1.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. nqxpack-0.1.1/.github/copybara/common-nqx.bara.sky +51 -0
  2. nqxpack-0.1.1/.github/copybara/common-template-copybara.bara.sky +65 -0
  3. nqxpack-0.1.1/.github/copybara/oss-nqx.bara.sky +5 -0
  4. nqxpack-0.1.1/.github/workflows/CI.yml +45 -0
  5. nqxpack-0.1.1/.github/workflows/copybara.yml +83 -0
  6. nqxpack-0.1.1/.github/workflows/formatting_check.yml +26 -0
  7. nqxpack-0.1.1/.github/workflows/publish.yml +38 -0
  8. nqxpack-0.1.1/.gitignore +147 -0
  9. nqxpack-0.1.1/.pre-commit-config.yaml +12 -0
  10. nqxpack-0.1.1/PKG-INFO +85 -0
  11. nqxpack-0.1.1/README.md +72 -0
  12. nqxpack-0.1.1/design/saving_mcstates.md +252 -0
  13. nqxpack-0.1.1/nqxpack/__init__.py +9 -0
  14. nqxpack-0.1.1/nqxpack/_src/__init__.py +1 -0
  15. nqxpack-0.1.1/nqxpack/_src/api.py +166 -0
  16. nqxpack-0.1.1/nqxpack/_src/contextmgr.py +108 -0
  17. nqxpack-0.1.1/nqxpack/_src/distributed.py +33 -0
  18. nqxpack-0.1.1/nqxpack/_src/errors.py +55 -0
  19. nqxpack-0.1.1/nqxpack/_src/io.py +129 -0
  20. nqxpack-0.1.1/nqxpack/_src/lib_v1/__init__.py +20 -0
  21. nqxpack-0.1.1/nqxpack/_src/lib_v1/asset_lib.py +178 -0
  22. nqxpack-0.1.1/nqxpack/_src/lib_v1/closure.py +69 -0
  23. nqxpack-0.1.1/nqxpack/_src/lib_v1/custom_types.py +163 -0
  24. nqxpack-0.1.1/nqxpack/_src/lib_v1/lib.py +246 -0
  25. nqxpack-0.1.1/nqxpack/_src/lib_v1/resolution.py +120 -0
  26. nqxpack-0.1.1/nqxpack/_src/lib_v1/versioninfo.py +1 -0
  27. nqxpack-0.1.1/nqxpack/_src/metadata/generate_metadata.py +50 -0
  28. nqxpack-0.1.1/nqxpack/_src/registry/__init__.py +11 -0
  29. nqxpack-0.1.1/nqxpack/_src/registry/flax.py +10 -0
  30. nqxpack-0.1.1/nqxpack/_src/registry/hydra.py +9 -0
  31. nqxpack-0.1.1/nqxpack/_src/registry/jax.py +33 -0
  32. nqxpack-0.1.1/nqxpack/_src/registry/netket.py +283 -0
  33. nqxpack-0.1.1/nqxpack/_src/registry/netket_operator.py +157 -0
  34. nqxpack-0.1.1/nqxpack/_src/registry/stdlib.py +33 -0
  35. nqxpack-0.1.1/nqxpack/_src/registry/versioninfo.py +1 -0
  36. nqxpack-0.1.1/nqxpack/registry.py +15 -0
  37. nqxpack-0.1.1/pyproject.toml +65 -0
  38. nqxpack-0.1.1/test/__init__.py +0 -0
  39. nqxpack-0.1.1/test/common/__init__.py +18 -0
  40. nqxpack-0.1.1/test/common/common.py +155 -0
  41. nqxpack-0.1.1/test/common/finite_diff.py +42 -0
  42. nqxpack-0.1.1/test/nqxpack/__init__.py +0 -0
  43. nqxpack-0.1.1/test/nqxpack/check.py +0 -0
  44. nqxpack-0.1.1/test/nqxpack/integration/__init__.py +0 -0
  45. nqxpack-0.1.1/test/nqxpack/integration/test_mcstate.py +73 -0
  46. nqxpack-0.1.1/test/nqxpack/integration/test_operators.py +46 -0
  47. nqxpack-0.1.1/test/nqxpack/test_hydratype_conversion.py +21 -0
@@ -0,0 +1,51 @@
1
+ push_commons = [
2
+ ".github/copybara/common-**",
3
+ "test/common/**",
4
+ ]
5
+ push_exclide_common = [
6
+ "**/_version.py",
7
+ "deepnets/**",
8
+ "ext/**",
9
+ "README.md",
10
+ "LICENSE",
11
+ "pyproject.toml",
12
+ ]
13
+ destination_common = [
14
+ ".github/copybara/common-**",
15
+ "test/common/**",
16
+ ]
17
+
18
+ copybara_transforms = [
19
+ core.move("packages/nqxpack", "nqxpack/"),
20
+ # core.move("test/nqxpack", "test/nqxpack"),
21
+ # core.move("test/common", "test/common"),
22
+ core.move("design/nqxpack", "design"),
23
+ # core.move("docs/packages/nqxpack", "docs/nqxpack"),
24
+ ]
25
+
26
+ os_config = {
27
+ "PUSH_TRANSFORMATIONS": copybara_transforms,
28
+ "PR_TRANSFORMATIONS": core.reverse(copybara_transforms),
29
+ "PUSH_INCLUDE": [
30
+ ".github/copybara/oss-nqx**",
31
+ "packages/nqxpack/**",
32
+ "test/nqxpack/**",
33
+ "design/nqxpack/**",
34
+ "examples/nqxpack/**",
35
+ "docs/packages/nqxpack/**",
36
+ ] + push_commons,
37
+ "PUSH_EXCLUDE": [
38
+
39
+ ] + push_exclide_common,
40
+ "DESTINATION_INCLUDE": [
41
+ ".github/copybara/oss-nqx**",
42
+ "nqxpack/**",
43
+ "test/nqxpack/**",
44
+ "design/**",
45
+ "examples/**",
46
+ "docs/nqxpack/**",
47
+ ] + destination_common,
48
+ "DESTINATION_EXCLUDE": [
49
+ "**/_version.py"
50
+ ]
51
+ }
@@ -0,0 +1,65 @@
1
+ # Variables
2
+ config = {
3
+ "SOT_REPO" : "git@github.com:<sot_repo>.git",
4
+ "SOT_BRANCH" : "<branch_to_sync>",
5
+ "DESTINATION_REPO" : "git@github.com:<destination_repo>.git",
6
+ "DESTINATION_BRANCH" : "<branch_to_sync>",
7
+ "COMMITTER" : "Quantum AI Lab Bot <quantumailabx@gmail.com>",
8
+ #"LOCAL_SOT" : "file:///usr/src/app",
9
+ }
10
+
11
+ def add_comment_with_shadow_pr(ctx):
12
+ ctx.console.info("Adding comment about internal PR to the original PR")
13
+ pr_number = ctx.revision.labels["GITHUB_PR_NUMBER"][0]
14
+ for effect in ctx.effects:
15
+ if effect.type == "CREATED" and effect.destination_ref != None and effect.destination_ref.url != None:
16
+ ctx.origin.post_issue_comment(int(pr_number), "Internal PR " + effect.destination_ref.url.replace("https://github.com/","").replace("/pull/","#") + "\nInternal message only. Nothing to see here, move along")
17
+
18
+ def pr_workflow(config, repo_config):
19
+ core.workflow(
20
+ name = "pr",
21
+ origin = git.github_pr_origin(
22
+ url = config["DESTINATION_REPO"],
23
+ branch = config["DESTINATION_BRANCH"],
24
+ ),
25
+ destination = git.github_pr_destination(
26
+ url = config["SOT_REPO"],
27
+ destination_ref = config["SOT_BRANCH"],
28
+ integrates = [],
29
+ title = "${GITHUB_PR_TITLE}",
30
+ update_description = True,
31
+ pr_branch = "imported/" + config["DESTINATION_REPO"].replace("git@github.com:", "").replace(".git", "") + "/${CONTEXT_REFERENCE}",
32
+ ),
33
+ destination_files = glob(repo_config["PUSH_INCLUDE"], exclude = repo_config["PUSH_EXCLUDE"]),
34
+ origin_files = glob(repo_config["DESTINATION_INCLUDE"], exclude = repo_config["DESTINATION_EXCLUDE"]),
35
+ authoring = authoring.pass_thru(default = config["COMMITTER"]),
36
+ mode = "CHANGE_REQUEST",
37
+ set_rev_id = False,
38
+ transformations = [
39
+ metadata.squash_notes(prefix='Import of ${GITHUB_PR_URL}\n\n**Original PR description:**\n\n${GITHUB_PR_BODY}\n\nImported changes:\n\n', oldest_first = True, show_author=False),
40
+ metadata.expose_label("GITHUB_PR_NUMBER", new_name = "Closes", separator = config["DESTINATION_REPO"].replace("git@github.com:", " ").replace(".git", "#")),
41
+ metadata.save_author("ORIGINAL_AUTHOR"),
42
+ ] + repo_config["PR_TRANSFORMATIONS"],
43
+ # after_workflow = [add_comment_with_shadow_pr],
44
+ )
45
+
46
+ def push_workflow(config, repo_config):
47
+ core.workflow(
48
+ name = "push",
49
+ origin = git.origin(
50
+ url = config["SOT_REPO"], # config["LOCAL_SOT"] if config["LOCAL_SOT"] else config["SOT_REPO"],
51
+ ref = config["SOT_BRANCH"],
52
+ ),
53
+ destination = git.github_destination(
54
+ url = config["DESTINATION_REPO"],
55
+ push = config["DESTINATION_BRANCH"],
56
+ ),
57
+ origin_files = glob(repo_config["PUSH_INCLUDE"], exclude = repo_config["PUSH_EXCLUDE"]),
58
+ destination_files=glob(repo_config["DESTINATION_INCLUDE"], exclude = repo_config["DESTINATION_EXCLUDE"]),
59
+ authoring = authoring.pass_thru(default = config["COMMITTER"]),
60
+ mode = "ITERATIVE",
61
+ transformations = [
62
+ metadata.restore_author("ORIGINAL_AUTHOR", search_all_changes = True),
63
+ metadata.expose_label("COPYBARA_INTEGRATE_REVIEW"),
64
+ ] + core.reverse(repo_config["PR_TRANSFORMATIONS"]),
65
+ )
@@ -0,0 +1,5 @@
1
+ load("app/.github/copybara/common-copybara", "config", "pr_workflow")
2
+
3
+ load("app/.github/copybara/common-nqx", "os_config")
4
+
5
+ pr_workflow(config, os_config)
@@ -0,0 +1,45 @@
1
+ name: CI
2
+
3
+ on:
4
+ pull_request:
5
+
6
+ push:
7
+ branches:
8
+ - master
9
+
10
+ concurrency:
11
+ group: ${{ github.workflow }}-${{ github.ref }}
12
+ cancel-in-progress: ${{ !contains(github.ref, 'master')}}
13
+
14
+
15
+ jobs:
16
+ test:
17
+ runs-on: ${{ matrix.os }}
18
+ strategy:
19
+ fail-fast: false
20
+ matrix:
21
+ include:
22
+ - os: ubuntu-latest
23
+ python-version: "3.12"
24
+
25
+ steps:
26
+ - uses: actions/checkout@v4
27
+ with:
28
+ submodules: true
29
+ fetch-tags: true
30
+
31
+ - name: Install the latest version of uv and set the python version
32
+ uses: astral-sh/setup-uv@v5
33
+ with:
34
+ python-version: ${{ matrix.python-version }}
35
+
36
+ - name: Tests
37
+ run: |
38
+ uv run pytest
39
+
40
+ # - name: Upload coverage reports to Codecov
41
+ # if: always()
42
+ # uses: codecov/codecov-action@v4
43
+ # with:
44
+ # token: ${{ secrets.CODECOV_TOKEN }} # required
45
+ # verbose: true # optional (default = false)
@@ -0,0 +1,83 @@
1
+ name: Copybara
2
+ on:
3
+ pull_request_target:
4
+ branches:
5
+ - "main"
6
+ types:
7
+ - reopened
8
+ - opened
9
+ - synchronize
10
+ - edited
11
+ env:
12
+ branch_to_sync: ${{ github.base_ref || github.ref_name}}
13
+ sot_repo: NeuralQXLab/netket_pro
14
+ destination_repo: NeuralQXLab/nqxpack
15
+ commiter_name: "Quantum AI Lab @ X - Automation Bot"
16
+ commiter_email: "quantumailabx@gmail.com"
17
+ copybara_config: ".github/copybara/oss-nqx.bara.sky"
18
+
19
+ jobs:
20
+ copybara-os-pr:
21
+ name: Mirror-pr
22
+ runs-on: ubuntu-latest
23
+ steps:
24
+ - uses: actions/checkout@v4
25
+
26
+ - name: Set up SSH key
27
+ run: |
28
+ mkdir -p ~/.ssh
29
+ echo "$SSH_PRIV_KEY" > ~/.ssh/id_rsa
30
+ chmod 600 ~/.ssh/id_rsa
31
+ ssh-keyscan -H github.com >> ~/.ssh/known_hosts
32
+ echo "https://x-access-token:${GH_TOKEN}@api.github.com" > ~/.git-credentials
33
+ env:
34
+ SSH_PRIV_KEY: ${{ secrets.SSH_GH_KEY }}
35
+ GH_TOKEN: ${{ secrets.GH_BOT_GITHUB_TOKEN }}
36
+
37
+ - name: Set up Git config
38
+ run: |
39
+ git config --global user.name "${{ env.commiter_name }}"
40
+ git config --global user.email "${{ env.commiter_email }}"
41
+
42
+ - name: Reopen referenced PRs
43
+ if: github.event.action == 'reopened'
44
+ run: |
45
+ PR_NUMBER="${{ github.event.pull_request.number }}"
46
+ CONTENT="$(gh pr view "$PR_NUMBER" --json "body,comments" -q '.body, .comments[].body' | tac)"
47
+
48
+ while read -r line; do
49
+ if [[ $line =~ (Closes|Internal\ PR)\ ([^/]+)/([^/]+)#([0-9]+) ]]; then
50
+ OWNER="${BASH_REMATCH[2]}"
51
+ REPO="${BASH_REMATCH[3]}"
52
+ OTHER_PR="${BASH_REMATCH[4]}"
53
+ echo "Reopening PR: $OWNER/$REPO/$OTHER_PR"
54
+ gh pr reopen --repo "$OWNER/$REPO" \
55
+ --comment "This pull request has been reopened because of reopening of https://github.com/${{ github.repository }}/pull/${PR_NUMBER}" \
56
+ $OTHER_PR
57
+ fi
58
+ done <<< "$CONTENT"
59
+ env:
60
+ GH_TOKEN: ${{ secrets.GH_BOT_GITHUB_TOKEN }}
61
+ - name: Generate copybara config file
62
+ run: |
63
+ sed -e 's#<branch_to_sync>#${{ env.branch_to_sync }}#' \
64
+ -e 's#<sot_repo>#${{ env.sot_repo }}#' \
65
+ -e 's#<destination_repo>#${{ env.destination_repo }}#' \
66
+ .github/copybara/common-template-copybara.bara.sky > .github/copybara/common-copybara.bara.sky
67
+
68
+ - name: Show config being used
69
+ run: |
70
+ cat .github/copybara/common-copybara.bara.sky
71
+ cat ${{ env.copybara_config }}
72
+
73
+ - uses: hazelcast/copybara-action@main
74
+ name: "Copybara PR #${{ github.event.pull_request.number }} -> ${{ env.branch_to_sync }}"
75
+ with:
76
+ ssh_key: ${{ secrets.SSH_GH_KEY }}
77
+ access_token: ${{ secrets.GH_BOT_GITHUB_TOKEN }}
78
+ sot_repo: ${{ env.sot_repo }}
79
+ sot_branch: ${{ env.branch_to_sync }}
80
+ destination_repo: ${{ env.destination_repo }}
81
+ copybara_options: "--nogit-destination-rebase"
82
+ committer: "${{ env.commiter_name }} <${{ env.commiter_email }}>"
83
+ custom_config: ${{ env.copybara_config }}
@@ -0,0 +1,26 @@
1
+ name: Checks
2
+
3
+ # Controls when the action will run. Triggers the workflow on push or pull request
4
+ # events but only for the master branch
5
+ on:
6
+ pull_request:
7
+
8
+ push:
9
+ branches:
10
+ - master
11
+
12
+ concurrency:
13
+ group: ${{ github.workflow }}-${{ github.ref }}
14
+
15
+
16
+ jobs:
17
+ pre-commit:
18
+ name: Format and Lint (pre-commit)
19
+ runs-on: ubuntu-latest
20
+ steps:
21
+ - uses: actions/checkout@v3
22
+ - uses: actions/setup-python@v3
23
+ with:
24
+ python-version: "3.12"
25
+
26
+ - uses: pre-commit/action@v3.0.1
@@ -0,0 +1,38 @@
1
+ # This workflows will upload a Python Package using Twine when a release is created
2
+ # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
3
+
4
+ name: Publish to PyPi
5
+
6
+ on:
7
+ release:
8
+ types: [published]
9
+
10
+ jobs:
11
+ deploy:
12
+ runs-on: ubuntu-latest
13
+ environment:
14
+ name: pypi
15
+ url: https://pypi.org/p/netket
16
+ permissions:
17
+ id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
18
+
19
+ steps:
20
+ - uses: actions/checkout@v4
21
+ with:
22
+ submodules: true
23
+
24
+ # make sure tags are fetched so we can get a version
25
+ - run: |
26
+ git fetch --prune --unshallow --tags
27
+
28
+ - name: Install the latest version of uv and set the python version
29
+ uses: astral-sh/setup-uv@v5
30
+ with:
31
+ python-version: "3.11"
32
+
33
+ - name: Build
34
+ run: |
35
+ uv build
36
+
37
+ - name: Publish package distributions to PyPI
38
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,147 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ uv.lock
6
+
7
+
8
+ # C extensions
9
+ *.so
10
+
11
+ # Distribution / packaging
12
+ netket_pro/_version.py
13
+ .Python
14
+ build/
15
+ develop-eggs/
16
+ dist/
17
+ downloads/
18
+ eggs/
19
+ .eggs/
20
+ lib/
21
+ lib64/
22
+ parts/
23
+ sdist/
24
+ var/
25
+ wheels/
26
+ pip-wheel-metadata/
27
+ share/python-wheels/
28
+ *.egg-info/
29
+ .installed.cfg
30
+ *.egg
31
+ MANIFEST
32
+
33
+ # PyInstaller
34
+ # Usually these files are written by a python script from a template
35
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
36
+ *.manifest
37
+ *.spec
38
+ .envrc
39
+
40
+ # Installer logs
41
+ pip-log.txt
42
+ pip-delete-this-directory.txt
43
+
44
+ # Unit test / coverage reports
45
+ htmlcov/
46
+ .tox/
47
+ .nox/
48
+ .coverage
49
+ .coverage.*
50
+ .cache
51
+ nosetests.xml
52
+ coverage.xml
53
+ *.cover
54
+ *.py,cover
55
+ .hypothesis/
56
+ .pytest_cache/
57
+
58
+ # Translations
59
+ *.mo
60
+ *.pot
61
+
62
+ # Django stuff:
63
+ *.log
64
+ local_settings.py
65
+ db.sqlite3
66
+ db.sqlite3-journal
67
+
68
+ # Flask stuff:
69
+ instance/
70
+ .webassets-cache
71
+
72
+ # Scrapy stuff:
73
+ .scrapy
74
+
75
+ # Sphinx documentation
76
+ docs/_build/
77
+
78
+ # PyBuilder
79
+ target/
80
+
81
+ # Jupyter Notebook
82
+ .ipynb_checkpoints
83
+ .virtual_ocuments
84
+
85
+ # IPython
86
+ profile_default/
87
+ ipython_config.py
88
+
89
+ # pyenv
90
+ .python-version
91
+
92
+ #jsons
93
+ *.json
94
+
95
+ # pipenv
96
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
97
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
98
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
99
+ # install all needed dependencies.
100
+ #Pipfile.lock
101
+
102
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow
103
+ __pypackages__/
104
+
105
+ # Celery stuff
106
+ celerybeat-schedule
107
+ celerybeat.pid
108
+
109
+ # SageMath parsed files
110
+ *.sage.py
111
+
112
+ # Environments
113
+ .env
114
+ .venv
115
+ env/
116
+ venv/
117
+ ENV/
118
+ env.bak/
119
+ venv.bak/
120
+
121
+ .DS_Store
122
+ # Spyder project settings
123
+ .spyderproject
124
+ .spyproject
125
+
126
+ # Rope project settings
127
+ .ropeproject
128
+
129
+ # mkdocs documentation
130
+ /site
131
+
132
+ # mypy
133
+ .mypy_cache/
134
+ .dmypy.json
135
+ dmypy.json
136
+ .ruff_cache
137
+ .lprof
138
+ .benchmarks
139
+
140
+ # Pyre type checker
141
+ .pyre/
142
+
143
+ # Website #
144
+ ###########
145
+ docs/_build
146
+ _generated
147
+ dist
@@ -0,0 +1,12 @@
1
+ repos:
2
+ - repo: https://github.com/astral-sh/ruff-pre-commit
3
+ # Ruff version.
4
+ rev: v0.9.4
5
+ hooks:
6
+ - id: ruff
7
+ args: [--fix, --exit-non-zero-on-fix]
8
+
9
+ - repo: https://github.com/psf/black
10
+ rev: 25.1.0
11
+ hooks:
12
+ - id: black
nqxpack-0.1.1/PKG-INFO ADDED
@@ -0,0 +1,85 @@
1
+ Metadata-Version: 2.4
2
+ Name: nqxpack
3
+ Version: 0.1.1
4
+ Summary: Library to serialise objects arising from scientific and Machine Learning libraries, including NetKet.
5
+ Author-email: Filippo Vicentini <filippovicentini@gmail.com>
6
+ Requires-Python: >=3.10
7
+ Requires-Dist: flax>=0.10.2
8
+ Requires-Dist: jax>=0.4.35
9
+ Requires-Dist: netket>=3.16.0.post2
10
+ Requires-Dist: numpy>=2
11
+ Requires-Dist: omegaconf>=2
12
+ Description-Content-Type: text/markdown
13
+
14
+ # NQXPack
15
+
16
+ A library to save and load objects coming from Scientific Machine Learning libraries, with a special attention to Neural Quantum States from NetKet.
17
+
18
+ Goals:
19
+ - Simple format, possible to hand-edit and inspect manually;
20
+ - Compatibility among Python version;
21
+ - Allows to load Neural Networks with a single ``load`` command;
22
+
23
+ ## Usage
24
+
25
+ Install with
26
+
27
+ ```bash
28
+ uv add git+https://github.com/NeuralQXLab/nqxpack.git
29
+ ```
30
+ or (but seriously, stop using pip and start using uv)
31
+ ```bash
32
+ pip install git+https://github.com/NeuralQXLab/nqxpack.git
33
+ ```
34
+
35
+
36
+ ### With `flax.linen`
37
+
38
+ Save a dictionary containing the model and the parameters.
39
+
40
+ ```python
41
+ import nqxpack
42
+ import jax
43
+ from flax import linen as nn
44
+
45
+ model = nn.Sequential(
46
+ nn.Dense(features=2),
47
+ nn.gelu,
48
+ nn.Dense(features=1),
49
+ jax.numpy.squeeze,
50
+ )
51
+
52
+ variables = model.init(jax.random.key(1), jax.numpy.ones((2,4)))
53
+
54
+ nqxpack.save({'model':model, 'variables':variables}, "mymodel.nk")
55
+
56
+ loaded_dict = nqxpack.load("mymodel.nk")
57
+ loaded_model, loaded_variables = loaded_dict['model'], loaded_dict['variables']
58
+ ```
59
+
60
+ ### With NetKet
61
+
62
+ ```python
63
+ import nqxpack
64
+ import netket as nk
65
+
66
+ hi = nk.hilbert.Spin(0.5, 10)
67
+ operator = nk.operator.spin.sigmax(nqs_state.hilbert, 1)
68
+
69
+ nqs_state = nk.vqs.MCState(nk.sampler.MetropolisLocal(hi), nk.models.RBM(alpha=4))
70
+ # print expectation value:
71
+ nqs_state.expect(operator)
72
+
73
+ nqxpack.save(nqs_state, "nqs_state.nk")
74
+ nqs_state_loaded = nqxpack.load("nqs_state.nk")
75
+
76
+ nqs_state_loaded.expect(operator)
77
+ ```
78
+
79
+ ## The format
80
+
81
+ The format is a single zip file. You can decompress it yourself and look into it.
82
+
83
+ ## Feedback required
84
+
85
+ If you use this library, please let us know of any issue you might find.
@@ -0,0 +1,72 @@
1
+ # NQXPack
2
+
3
+ A library to save and load objects coming from Scientific Machine Learning libraries, with a special attention to Neural Quantum States from NetKet.
4
+
5
+ Goals:
6
+ - Simple format, possible to hand-edit and inspect manually;
7
+ - Compatibility among Python version;
8
+ - Allows to load Neural Networks with a single ``load`` command;
9
+
10
+ ## Usage
11
+
12
+ Install with
13
+
14
+ ```bash
15
+ uv add git+https://github.com/NeuralQXLab/nqxpack.git
16
+ ```
17
+ or (but seriously, stop using pip and start using uv)
18
+ ```bash
19
+ pip install git+https://github.com/NeuralQXLab/nqxpack.git
20
+ ```
21
+
22
+
23
+ ### With `flax.linen`
24
+
25
+ Save a dictionary containing the model and the parameters.
26
+
27
+ ```python
28
+ import nqxpack
29
+ import jax
30
+ from flax import linen as nn
31
+
32
+ model = nn.Sequential(
33
+ nn.Dense(features=2),
34
+ nn.gelu,
35
+ nn.Dense(features=1),
36
+ jax.numpy.squeeze,
37
+ )
38
+
39
+ variables = model.init(jax.random.key(1), jax.numpy.ones((2,4)))
40
+
41
+ nqxpack.save({'model':model, 'variables':variables}, "mymodel.nk")
42
+
43
+ loaded_dict = nqxpack.load("mymodel.nk")
44
+ loaded_model, loaded_variables = loaded_dict['model'], loaded_dict['variables']
45
+ ```
46
+
47
+ ### With NetKet
48
+
49
+ ```python
50
+ import nqxpack
51
+ import netket as nk
52
+
53
+ hi = nk.hilbert.Spin(0.5, 10)
54
+ operator = nk.operator.spin.sigmax(nqs_state.hilbert, 1)
55
+
56
+ nqs_state = nk.vqs.MCState(nk.sampler.MetropolisLocal(hi), nk.models.RBM(alpha=4))
57
+ # print expectation value:
58
+ nqs_state.expect(operator)
59
+
60
+ nqxpack.save(nqs_state, "nqs_state.nk")
61
+ nqs_state_loaded = nqxpack.load("nqs_state.nk")
62
+
63
+ nqs_state_loaded.expect(operator)
64
+ ```
65
+
66
+ ## The format
67
+
68
+ The format is a single zip file. You can decompress it yourself and look into it.
69
+
70
+ ## Feedback required
71
+
72
+ If you use this library, please let us know of any issue you might find.