Skip to content

Commit 3e22249

Browse files
author
jurajHasik
committed
Merge branch 'master' of github.com:jurajHasik/tn-torch_dev
2 parents b25f49b + e17857f commit 3e22249

File tree

149 files changed

+110771
-1974
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

149 files changed

+110771
-1974
lines changed

.github/workflows/main.yml

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@ jobs:
2121
runs-on: ubuntu-latest
2222
strategy:
2323
matrix:
24-
python_v: [3.7, 3.8, 3.9]
25-
pytorch_v: [1.11.0]
26-
fail-fast: false
24+
python_v: [ "3.8", "3.9", "3.10", "3.11"]
25+
pytorch_v: [2.0]
26+
fail-fast: false
2727

2828
# Steps represent a sequence of tasks that will be executed as part of the job
2929
steps:
@@ -47,7 +47,7 @@ jobs:
4747
then
4848
$CONDA/bin/conda install -c pytorch-lts -c anaconda pytorch cpuonly scipy pytest
4949
else
50-
$CONDA/bin/conda install -c pytorch -c anaconda pytorch==${{ matrix.pytorch_v }} cpuonly scipy pytest
50+
$CONDA/bin/conda install -c pytorch -c anaconda -c conda-forge pytorch==${{ matrix.pytorch_v }} cpuonly scipy pytest opt_einsum
5151
fi
5252
git submodule update --init --recursive
5353
@@ -85,6 +85,7 @@ jobs:
8585
run: |
8686
$CONDA/bin/python -m pytest examples/ladders/abelian/ctmrg_*.py
8787
$CONDA/bin/python -m pytest examples/ladders/abelian/optim_*.py
88+
$CONDA/bin/python -m pytest examples/ladders/abelian/SU_ladders_u1.py
8889
8990
- name: test kagome
9091
run: |

.gitmodules

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
1-
[submodule "yast"]
2-
path = yast
3-
url = https://gitlab.com/marekrams/yast.git
1+
[submodule "yastn"]
2+
path = yastn
3+
url = https://github.com/yastn/yastn.git

README.md

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ by the corner-transfer matrix (CTM) algorithm. Afterwards, the gradients are com
1111
automatic differentiation (AD).
1212

1313
#### Now supporting
14-
* **abelian symmetries, with implementation powered by [YAST](https://gitlab.com/marekrams/yast)**
14+
* **abelian symmetries, with implementation power by [YASTN](https://github.com/yastn/yastn)**
1515

1616
Allows definition of abelian-symmetric iPEPS in terms of block-sparse tensors, computing their
1717
symmetric environments, and their optimization with gradient-based methods.
@@ -129,10 +129,12 @@ python examples/j1j2/abelian/ctmrg_j1j2_u1.py --tiling BIPARTITE --chi 48 --j2 0
129129

130130
#### Dependencies
131131
- PyTorch 1.11+ (see https://pytorch.org/)
132-
- (optional) YAST (see https://gitlab.com/marekrams/yast)
132+
- (optional) YASTN (see https://github.com/yastn/yastn)
133133
- (optional) scipy 1.3.+
134+
- (optional) opt_einsum
135+
- (optional) ArrayFire (see https://github.com/arrayfire/arrayfire)
134136

135-
YAST is linked to **peps-torch** as a [git submodule](https://git-scm.com/book/en/v2/Git-Tools-Submodules).
137+
YASTN is linked to **peps-torch** as a [git submodule](https://git-scm.com/book/en/v2/Git-Tools-Submodules).
136138
To obtain it, you can use git:
137139

138140
`git submodule update --init --recursive`

benchmarks/conftest.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import pytest
2+
import context
3+
import config as cfg
4+
5+
def pytest_addoption(parser):
6+
parser.addoption(
7+
"--device",
8+
default="cpu",
9+
help="choose device",
10+
)
11+
12+
def pytest_configure(config):
13+
device= config.getoption("--device")
14+
cfg.global_args.device= device

benchmarks/context.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
import os
2+
import sys
3+
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))

benchmarks/test_energy_j1j2trgl.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import context
2+
import pytest
3+
import torch
4+
import config as cfg
5+
from ipeps.ipeps import IPEPS
6+
from ctm.generic.env import ENV, init_random
7+
from models.spin_triangular import J1J2J4_1SITE, eval_nn_per_site, eval_nnn_per_site, eval_nn_and_chirality_per_site
8+
9+
import logging
10+
logging.basicConfig(filename=f"{__file__}.log", filemode='w', level=logging.INFO)
11+
12+
test_dims=[(3,27), (3,54), (4,32)]
13+
14+
15+
@pytest.mark.parametrize("dims",test_dims)
16+
@pytest.mark.parametrize("unroll",[True,False])
17+
def test_profile_j1j2_loop_oe_semimanual(dims, unroll, benchmark):
18+
D,X= dims
19+
20+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
21+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
22+
env= ENV(X, state)
23+
init_random(env)
24+
25+
model= J1J2J4_1SITE(phys_dim=2, j1=1.0, j2=1.0, j4=0, jchi=0, global_args=cfg.global_args)
26+
27+
def test_f():
28+
nn_h_v,nn_diag= eval_nn_per_site((0,0),state,env,model.R,model.R@model.R,model.SS,model.SS)
29+
nnn= eval_nnn_per_site((0,0),state,env,None,None,model.SS,looped=unroll,use_checkpoint=False)
30+
31+
benchmark.pedantic(test_f, args=(),\
32+
iterations=1, rounds=2, warmup_rounds=1)
33+
34+
@pytest.mark.parametrize("dims",test_dims)
35+
@pytest.mark.parametrize("unroll",[True,False])
36+
def test_profile_j1j2jX_loop_oe_semimanual(dims, unroll, benchmark):
37+
D,X= dims
38+
39+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
40+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
41+
env= ENV(X, state)
42+
init_random(env)
43+
44+
model= J1J2J4_1SITE(phys_dim=2, j1=1.0, j2=1.0, j4=0, jchi=0, global_args=cfg.global_args)
45+
46+
def test_f():
47+
nnn= eval_nnn_per_site((0,0),state,env,None,None,model.SS,looped=unroll,use_checkpoint=False)
48+
nn_h_v,nn_diag,chi= eval_nn_and_chirality_per_site((0,0),state,env,\
49+
model.R,model.R@model.R,model.SS,model.SS,model.h_chi,looped=unroll,use_checkpoint=False)
50+
51+
benchmark.pedantic(test_f, args=(),\
52+
iterations=1, rounds=2, warmup_rounds=1)
53+
54+
@pytest.mark.parametrize("dims",test_dims)
55+
@pytest.mark.parametrize("unroll",[True,False])
56+
def test_profile_rdm2x3_loop_manual(dims, unroll, benchmark):
57+
D,X= dims
58+
59+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
60+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
61+
env= ENV(X, state)
62+
init_random(env)
63+
64+
model= J1J2J4_1SITE(phys_dim=2, j1=1.0, j2=1.0, j4=0, jchi=0, global_args=cfg.global_args)
65+
66+
benchmark.pedantic(model.energy_1x3, args=(state,env,-1,unroll,\
67+
cfg.ctm_args,cfg.global_args),\
68+
iterations=1, rounds=2, warmup_rounds=1)

benchmarks/test_rdm_2x2

Whitespace-only changes.

benchmarks/test_rdm_2x3.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import context
2+
import pytest
3+
import torch
4+
import config as cfg
5+
from ipeps.ipeps import IPEPS
6+
from ctm.generic.env import ENV, init_random
7+
from ctm.generic import rdm_mc
8+
9+
import logging
10+
logging.basicConfig(filename=f"{__file__}.log", filemode='w', level=logging.INFO)
11+
12+
test_dims=[(3,27), (3,54), (4,32)]
13+
14+
15+
@pytest.mark.parametrize("dims",test_dims)
16+
@pytest.mark.parametrize("open_inds",[[0,1,2,3,4,5]])
17+
@pytest.mark.parametrize("unroll",[True,False])
18+
def test_profile_rdm2x3_loop_oe(dims, open_inds, unroll, benchmark):
19+
D,X= dims
20+
21+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
22+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
23+
env= ENV(X, state)
24+
init_random(env)
25+
26+
print(f"{dims} {unroll}")
27+
benchmark.pedantic(rdm_mc.rdm2x3_loop_oe, args=((0,0), state, env, open_inds, unroll),\
28+
iterations=1, rounds=2, warmup_rounds=1)
29+
30+
@pytest.mark.parametrize("dims",test_dims)
31+
@pytest.mark.parametrize("open_inds",[[0,1,2,3,4,5]])
32+
@pytest.mark.parametrize("unroll",[True,False])
33+
def test_profile_rdm2x3_loop_oe_semimanual(dims, open_inds, unroll, benchmark):
34+
D,X= dims
35+
36+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
37+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
38+
env= ENV(X, state)
39+
init_random(env)
40+
41+
print(f"{dims} {unroll}")
42+
benchmark.pedantic(rdm_mc.rdm2x3_loop_oe_semimanual, args=((0,0), state, env, open_inds, unroll),\
43+
iterations=1, rounds=2, warmup_rounds=1)
44+
45+
46+
@pytest.mark.parametrize("dims",test_dims)
47+
def test_profile_rdm2x3_loop_manual(dims, benchmark):
48+
D,X= dims
49+
50+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
51+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
52+
env= ENV(X, state)
53+
init_random(env)
54+
55+
benchmark.pedantic(rdm_mc.rdm2x3_loop, args=((0,0), state, env),\
56+
iterations=1, rounds=2, warmup_rounds=1)

benchmarks/test_rdm_2x3_trglnnn.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import context
2+
import pytest
3+
import torch
4+
import config as cfg
5+
from ipeps.ipeps import IPEPS
6+
from ctm.generic.env import ENV, init_random
7+
from ctm.generic import rdm_mc
8+
9+
test_dims=[(3,27), (3,54), (4,32), (4,64)]
10+
11+
@pytest.mark.parametrize("dims",test_dims)
12+
@pytest.mark.parametrize("open_inds",[[2,3]])
13+
@pytest.mark.parametrize("unroll",[True,False])
14+
def test_profile_rdm2x3_loop_oe(dims, open_inds, unroll, benchmark):
15+
D,X= dims
16+
17+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
18+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
19+
env= ENV(X, state)
20+
init_random(env)
21+
22+
print(f"{dims} {unroll}")
23+
benchmark.pedantic(rdm_mc.rdm2x3_loop_oe, args=((0,0), state, env, open_inds, unroll),\
24+
iterations=1, rounds=2, warmup_rounds=1)
25+
26+
@pytest.mark.parametrize("dims",test_dims)
27+
@pytest.mark.parametrize("open_inds",[[2,3]])
28+
@pytest.mark.parametrize("unroll",[True,False])
29+
def test_profile_rdm2x3_loop_oe_semimanual(dims, open_inds, unroll, benchmark):
30+
D,X= dims
31+
32+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
33+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
34+
env= ENV(X, state)
35+
init_random(env)
36+
37+
print(f"{dims} {unroll}")
38+
benchmark.pedantic(rdm_mc.rdm2x3_loop_oe_semimanual, args=((0,0), state, env, open_inds, unroll),\
39+
iterations=1, rounds=2, warmup_rounds=1)

benchmarks/test_rdm_2x3_trglringex.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
import context
2+
import pytest
3+
import torch
4+
import functools
5+
from time import perf_counter
6+
import config as cfg
7+
from ipeps.ipeps import IPEPS
8+
from ctm.generic.env import ENV, init_random
9+
from ctm.generic import rdm_mc
10+
11+
import logging
12+
logging.basicConfig(filename=f"{__file__}.log", filemode='w', level=logging.INFO)
13+
14+
test_dims=[(3,27), (3,54), (4,32), (4,64)]
15+
16+
def optional_cuda_measure(tag):
17+
def _inner_optional_cuda_measure(f):
18+
@functools.wraps(f)
19+
def _wrap(*args,**kwargs):
20+
if not cfg.global_args.device=='cpu': torch.cuda.synchronize()
21+
t0= perf_counter()
22+
res= f(*args,**kwargs)
23+
if not cfg.global_args.device=='cpu': torch.cuda.synchronize()
24+
t1= perf_counter()
25+
logging.info(f"{tag} {t1-t0} [s]")
26+
return res
27+
return _wrap
28+
return _inner_optional_cuda_measure
29+
30+
@pytest.mark.parametrize("dims",test_dims)
31+
@pytest.mark.parametrize("open_inds",[[1,2,3,4]])
32+
@pytest.mark.parametrize("unroll",[True,False])
33+
def test_profile_rdm2x3_loop_oe(dims, open_inds, unroll, benchmark):
34+
D,X= dims
35+
36+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
37+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
38+
env= ENV(X, state)
39+
init_random(env)
40+
41+
@optional_cuda_measure(f"rdm2x3_loop_oe{dims} {open_inds} {unroll}")
42+
def test_f():
43+
rdm_mc.rdm2x3_loop_oe((0,0), state, env, open_inds, unroll)
44+
45+
print(f"{dims} {unroll}")
46+
benchmark.pedantic(test_f, args=(),\
47+
iterations=1, rounds=2, warmup_rounds=1)
48+
49+
@pytest.mark.parametrize("dims",test_dims)
50+
@pytest.mark.parametrize("open_inds",[[1,2,3,4]])
51+
@pytest.mark.parametrize("unroll",[True,False])
52+
def test_profile_rdm2x3_loop_oe_semimanual(dims, open_inds, unroll, benchmark):
53+
D,X= dims
54+
55+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
56+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
57+
env= ENV(X, state)
58+
init_random(env)
59+
60+
@optional_cuda_measure(f"rdm2x3_loop_oe_semimanual{dims} {open_inds} {unroll}")
61+
def test_f():
62+
rdm_mc.rdm2x3_loop_oe_semimanual((0,0), state, env, open_inds, unroll)
63+
64+
print(f"{dims} {unroll}")
65+
benchmark.pedantic(test_f, args=(),\
66+
iterations=1, rounds=2, warmup_rounds=1)
67+
68+
69+
@pytest.mark.parametrize("dims",test_dims)
70+
def test_profile_rdm2x3_loop_trglringex_manual(dims, benchmark):
71+
D,X= dims
72+
73+
state= IPEPS({(0,0): torch.rand((2,)+(D,)*4,\
74+
dtype=cfg.global_args.torch_dtype,device=cfg.global_args.device)-0.5}, lX=1, lY=1)
75+
env= ENV(X, state)
76+
init_random(env)
77+
78+
@optional_cuda_measure(f"rdm2x3_loop_trglringex_manual_{dims}")
79+
def test_f():
80+
rdm_mc.rdm2x3_loop_trglringex_manual((0,0), state, env)
81+
82+
83+
benchmark.pedantic(test_f, args=(),\
84+
iterations=1, rounds=2, warmup_rounds=1)

config.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,8 @@ class PEPSARGS():
199199
def __init__(self):
200200
self.build_dl= True
201201
self.build_dl_open= False
202+
self.quasi_gauge_max_iter= 10**6
203+
self.quasi_gauge_tol= 1.0e-8
202204

203205
def __str__(self):
204206
res=type(self).__name__+"\n"
@@ -256,6 +258,8 @@ class CTMARGS():
256258
:vartype projector_eps_multiplet: float
257259
:ivar projector_multiplet_abstol: absolute threshold for spectral values to be considered in multiplets
258260
:vartype projector_multiplet_abstol: float
261+
:ivar radomize_ctm_move_sequence: If ``True``, then ``ctm_move_sequence`` is randomized in each optimization step
262+
:vartype radomize_ctm_move_sequence: bool
259263
:ivar ctm_move_sequence: sequence of directional moves within single CTM iteration. The possible
260264
directions are encoded as tuples(int,int)
261265
@@ -329,6 +333,7 @@ def __init__(self):
329333
self.projector_multiplet_abstol = 1.0e-14
330334
self.ad_decomp_reg= 1.0e-12
331335
self.ctm_move_sequence = [(0,-1), (-1,0), (0,1), (1,0)]
336+
self.randomize_ctm_move_sequence = False
332337
self.ctm_force_dl = False
333338
self.ctm_logging = False
334339
self.verbosity_initialization = 0
@@ -342,6 +347,7 @@ def __init__(self):
342347
self.fwd_checkpoint_projectors = False
343348
self.fwd_checkpoint_absorb = False
344349
self.fwd_checkpoint_move = False
350+
self.fwd_checkpoint_loop_rdm = False
345351

346352
def __str__(self):
347353
res=type(self).__name__+"\n"
@@ -419,6 +425,7 @@ def __init__(self):
419425
self.tolerance_grad= 1e-5
420426
self.tolerance_change= 1e-9
421427
self.opt_ctm_reinit= True
428+
self.env_sens_scale= 10.0
422429
self.line_search= "default"
423430
self.line_search_ctm_reinit= True
424431
self.line_search_svd_method= 'DEFAULT'

0 commit comments

Comments
 (0)