feat(templates): Implement unified lattice module and tests

[Stellogic]

Hi @yutuer21,

Following your guidance and the OSPP project requirements ("Object-Oriented Lattice Generation"), I've prepared this initial pull request to add the unified lattice module. I would be grateful for your review and feedback.

Main Changes in this PR:

• AbstractLattice: An abstract base class is proposed to define the common API.
• Concrete Implementations: CustomizeLattice and TILattice (with SquareLattice, HoneycombLattice, TriangularLattice examples) are provided as initial implementations.
• Core Functionality: Basic methods like get_site_info(), get_neighbors(), and .show() visualization are included.
• Testing: A corresponding test suite (tests/test_lattice.py) has been added to cover the core functions and edge cases, achieving ~94% coverage on lattice.py.
• Documentation: CHANGELOG.md has been updated.

This is the first version based on my understanding of the requirements, and I'm sure there is still much room for improvement. Thank you very much for your valuable time and guidance. I look forward to your feedback!

[Stellogic]

This update primarily refactors the plotting logic within the show() method to resolve a persistent mypy [arg-type] error related to ax.text().

Instead of simply suppressing the error with # type: ignore, I opted for a more robust solution by using isinstance(ax, Axes3D) to perform a direct, type-safe check on the Matplotlib axes object. This approach makes the code's intent clearer to static analysis tools and fixes the type ambiguity cleanly.

I thought this time my code may pass the mypy test .

[Stellogic]

Following the guidance of pylint , I fix my code .
I thought this time my code may pass the ci.

[refraction-ray]

some more TILattice,
1D: chain, chain with AB sublattice (ABABAB...),
2D: rectangular, checkboard (i.e. square lattice with AB sublattice), Lieb lattice, Kagome lattice
3D: cubic lattice

[refraction-ray]

please also modify init.py in templates module accordingly

[Stellogic]

Hi @refraction-ray,

Thank you again for your invaluable guidance and detailed review. I have now pushed the updates, incorporating all of your feedback.

Here is a summary of the main changes:

• Logging & On-Demand Computation: Replaced all print calls with logging and implemented on-demand neighbor calculations.
• New TILattice Classes: Added and tested all the suggested lattice implementations (Chain, Kagome, Lieb, etc.).
• Code Style & Typing:
  • Removed type hints from test function arguments as requested.
  • In response to your question about the mypy error, I have added a detailed comment to the source code explaining the purpose of the if TYPE_CHECKING: block for handling the optional matplotlib dependency.
• Docstring Format Correction: I have carefully revised all docstrings in lattice.py to match the project's required style.

[refraction-ray]

Thanks for the nice contribution. In general, the code looks very good.

The addition should be a more robust and error-free neighbor searching implementation with careful consideration on PBC and large max_k

current failed case for neighbor search:

sq = lattice.SquareLattice(size=(4,4), pbc=True)
len(sq.get_neighbor_pairs(k=3))
# 0

Considering that the typical total size of the lattice used in tensorcircuit is 10-1000, it may be straightforward and robust to directly precompute the all-to-all full distance matrix with the real space distance in terms of pbc (see Minimum Image Convention method). There might be corner case with mixed boundary conditions and skewed lattice vector not align with pbc dimensions, let us just ignore such pathological cases.

[Stellogic]

Hi @refraction-ray,
Thanks again for your excellent guidance. I've pushed an update that I hope addresses your feedback.
Here's a quick summary of the changes:
The neighbor search algorithm has been completely refactored. It's now more robust for all boundary conditions (PBC, OBC, mixed) and significantly more performant, especially for periodic lattices. This also fixes the bug you found for k=3.
I've added the new tests you suggested for larger lattices, longer-range k values, and mixed-BCs.
CustomizeLattice now uses a KDTree-based method for better efficiency.
I'm eager to hear your thoughts. Thank you for your time and help!

[Stellogic]

@refraction-ray Fixed. I've pushed the commit to add the missing docstring. Thanks for the review.

[refraction-ray]

Very impressive and responsive work!

1. I suggest to store the all-to-all distance matrix into the lattice class when computing it, as the info is very frequently used, say, to implement an all-to-all long range interaction based on distance
2. for KDTree neighbor search, I suggest you to do some benchmark on efficiency (running time) compare to the baseline neighbor search method for number of sites ranging from 10 to 2000, to see whether there is any practical performance advantage in this regime. If necessary, the neighbor search engine (plain all-to-all distance vs. KDtree based index) should be exposed as one switch arguments for CustomizeLattice.
3. To better interpolate different lattice class, I suggest to add the following APIs for CustomizeLattice, so that we can distort from some TI lattice:

sq = SquareLattice([3, 3])
cl = CustomizeLattice.from_lattice(sq)
cl.add_sites(identifies, coordinates)
# error when identifies is duplicated internally
cl.remove_sites(identifies)

Again, the contribution overall looks very nice and impressive to me. (The current CI error is irrelevant from your commits, just ignore it)

[Stellogic]

Hi @refraction-ray
Thanks for the great suggestions! I've now pushed all the updates.
I've implemented the dynamic modification APIs (from_lattice, add_sites, remove_sites) for CustomizeLattice.
I've fixed the default show() behavior to avoid warnings.

I ran a performance benchmark for the neighbor finding methods. The data shows that the KDTree-based method is likely consistently faster. Based on this, I've kept the implementation simple, using KDTree by default without needing a user-facing switch. The benchmark script is included in templates/ .benchmarks/.

New tests for all new features have been added.

[refraction-ray]

while assisted with AI is okay, please ensure the change is necessary and minimal, (say obviously the frequent docstring & comment & variable name change is meaningless and gives negative contribution to code review as the diff is large). The latest commit, while indeed contains the distance_matrix cache for all cases, but introduces more unnecessary and even suboptimal changes to irrelevant parts of the module, please revert the unnecessary changes or explicit state why the changes matter

[Stellogic]

@refraction-ray
My apologies. My process before updating a PR has been to let an AI review it first for any potential issues. I would then follow its suggestions to modify things like variable names, logic, and docstrings. Going forward, I will be more careful in reviewing the AI's suggestions to ensure they are reasonable. Again, I am very sorry.

[refraction-ray]

@Stellogic , No need to apologize. It is super okay to use AI, actually I strongly recommend you guys to intensively use AI for coding and development tasks. Just pay attention to the points I mention, as a large bulk of unnecessary changes will make the code review more difficult. AI can write the code for the most parts and everyone use it. It is a good thing, but you should be the first responsible reviewer for what AI generates and changes :)

[Stellogic]

Hi @refraction-ray,

Following your guidance, I have force-pushed an update. I've reset the branch back to commit 9ef578d and built a single, clean commit on top of it containing only the necessary fixes to address the test failures.

This new commit includes three targeted fixes:

Fixed Distance Matrix Caching for TILattice: In the fully periodic case, _build_neighbors now reconstructs and caches the full distance matrix by leveraging the existing ref_dist_matrix_sq and translational symmetry. This resolves all failures in TestDistanceMatrix.

Fixed Neighbor Finding for Duplicate Coordinates in
CustomizeLattice: The logic in _build_neighbors has been refined. It now correctly handles co-located sites by ensuring all points at distance ~0 are excluded from the initial k=1 neighbor shell, which fixes the test_lattice_with_duplicate_coordinates failure.

Improved Numerical Stability in _identify_distance_shells: The logic for filtering zero-distance points has been decoupled from the tolerance-based shell identification. It now robustly filters out self-loops using a fixed small threshold (1e-12), resolving the test_neighbor_shells_with_tiny_separation failure.

[refraction-ray]

tol**2 instead of 1e-12?

[Stellogic]

I've looked into this, and it seems that a direct replacement with tol2 causes the test_neighbor_shells_with_tiny_separation test to fail.
The root cause seemms that tol is treated as a linear distance tolerance, but using tol2 in the squared-distance comparison logic makes the shell-merging behavior overly sensitive and breaks the test's intent.

[refraction-ray]

@Stellogic it is weird, as tol**2 = 1e-6**2 is just 1e-12? what makes the difference? otherwise, what is the point for the tol argument?

[refraction-ray]

besides, if the 1e-12 has no dependce on tol, how the test you mention can differentiate the two cases with different tol

[Stellogic]

Hi @refraction-ray,

I've pushed the new commit addressing your feedback.

The indentation issue with the _distance_matrix assignment has been fixed as you suggested.

But I've kept the 1e-12 implementation. The reason is that changing it to tol**2 causes the numerical stability test (test_neighbor_shells_with_tiny_separation) to fail.

[refraction-ray]

the CI error seems to be irrelevant for this PR. However, I am still confused on the success of test_neighbor_shells_with_tiny_separation function

[refraction-ray]

should this be outside the for as it doesn't depend on i?

[refraction-ray]

this should also be outside the for and created only once?

[refraction-ray]

max_k or len(dist_shells_sq)?

[refraction-ray]

sq = lattice.HoneycombLattice(size=(4,5), pbc=True, precompute_neighbors=1)
sq._neighbor_maps # now a full map instead of max_k = 1

[refraction-ray]

currently both pbc and obc in TILattice automatically build the whole neighbor maps with invalid max_k or precompute_neighbors

[refraction-ray]

my suggestion is to not use the special implementation of fully PBC case, as it is much slower and greatly complicate the implementation

[refraction-ray]

the mixed boundary condition universal implementation is efficient enough, just deleting the special tailored implementation branch for fully PBC case, and everything would be nicer and easier

[refraction-ray]

overall,

1. the 1e-12 issue is still a mystery to me
2. delete fully PBC part of the neighbor building and distance matrix construct
3. max_k should work for TILattice and CustomizeLattice to control the size of _neighbor_maps, currently both fail and always generate a full neighbor map

[Stellogic]

Hi@refraction-ray！

To clarify the point about 1e-12: it is indeed intended to handle floating-point precision issues by filtering out squared distances that are effectively zero. This is a fixed threshold and is independent of the tol parameter, which is used for separating distinct neighbor shells. The test_neighbor_shells_with_tiny_separation test, which passes various tol values, demonstrates this intended separation of concerns.

I've also strengthened the test suite with two additions:
A new test suite to verify that the max_k logic works as expected for both pre-computation and on-demand calculation.

A performance regression test to ensure we don't encounter performance degradation in the future.

[refraction-ray]

looks good to me now, thanks for the nice contribution! will merge soon

[refraction-ray]

@all-contributors please add @Stellogic for code, examples, tests

[allcontributors]

@refraction-ray

I've put up a pull request to add @Stellogic! 🎉

[refraction-ray]

Apart from some more tutorials and examples, the next key milestone could be a machine learning backend agnostic implementation for lattice module, compared to the numpy backend implementation now. A backend agnostic implementation using tensorcircuit.backend API would enable automatic differentiation, vectorized map etc. for this lattice module. A typical example is structure optimization, where the energy computed from the lattice is minimized via gradient descent with tunable parameters the lattice structure parameters.

[Stellogic]

Hi @refraction-ray,
Thank you so much for your guidance and for merging the PR! It has been a fantastic learning experience for me.
You are absolutely right about the need for tutorials and examples to make the new module accessible. I will start working on that and will submit it as a new PR to complete this contribution.
As for the excellent suggestion of a backend-agnostic implementation, I will definitely keep that in mind for a future contribution when I have more availability.
Thanks again for your mentorship!