update docs

[chaoming0625]

This pull request includes several changes to improve code readability, functionality, and performance. The most important changes include modifications to the pinnx/_trainer.py file to add new functionality for measuring training step compile and execution times, and improvements to the docs/unit-examples-forward/Beltrami_flow.py and docs/unit-examples-forward/heat.ipynb files for better code structure and clarity.

Functionality Enhancements:

• pinnx/_trainer.py: Added functionality to measure training step compile time and execution time by introducing new parameters measture_train_step_compile_time and measture_train_step_time in the compile and train methods respectively. [1] [2]

Code Structure Improvements:

• docs/unit-examples-forward/Beltrami_flow.py: Improved readability by reformatting the dictionary r in the icbc_cond_func function.
• docs/unit-examples-forward/heat.ipynb: Removed an unnecessary import statement and fixed a lambda function to improve code clarity. [1] [2]

Import Adjustments:

• pinnx/_trainer.py: Added import time to support the new timing functionality.
• pinnx/geometry/geometry_2d.py: Removed an unused import jax.numpy and added from pinnx import utils for better code organization.

Summary by Sourcery

Add functionality to measure training step compile and execution times, and improve code readability in documentation examples.

New Features:

• Measure training step compile and execution times.

Tests:

• Updated documentation examples for better code structure and clarity.

[sourcery-ai]

Reviewer's Guide by Sourcery

This pull request enhances the pinnx library by adding functionality to measure training step compile and execution times, along with improvements to code structure and clarity in example files. The _trainer.py file is modified to include new parameters for measuring these times, while the Beltrami_flow.py and heat.ipynb files in the documentation are updated for better readability and code organization.

Sequence diagram for training step timing measurement

sequenceDiagram
    participant Client
    participant Trainer
    participant TrainState

    Client->>Trainer: compile(measture_train_step_compile_time=true)
    activate Trainer
    Note right of Trainer: Start compile time measurement
    Trainer->>Trainer: _compile_training_step()
    Trainer->>TrainState: set_data_train()
    Trainer->>Trainer: fn_train_step.compile()
    Note right of Trainer: End compile time measurement
    Trainer-->>Client: return compile_time
    deactivate Trainer

    Client->>Trainer: train(measture_train_step_time=true)
    activate Trainer
    Note right of Trainer: Start execution time measurement
    Trainer->>Trainer: _train()
    Note right of Trainer: Training iterations
    Note right of Trainer: End execution time measurement
    Trainer-->>Client: return execution_time
    deactivate Trainer

Loading

Class diagram for Trainer modifications

classDiagram
    class Trainer {
        +compile(optimizer, metrics, measture_train_step_compile_time)
        +train(iterations, display_every, batch_size, callbacks, model_restore_path, model_save_path, measture_train_step_time)
        -_compile_training_step(batch_size)
        -_train(iterations, display_every, batch_size, callbacks)
    }
    note for Trainer "Added timing measurement capabilities"

Loading

File-Level Changes

Change	Details	Files
Added functionality to measure training step compile and execution times.	Introduced measture_train_step_compile_time parameter in the compile method. Introduced measture_train_step_time parameter in the train method. Added a _compile_training_step method to compile the training step function.	pinnx/_trainer.py
Improved code readability in the Beltrami flow example.	Reformatted the dictionary r in the icbc_cond_func function for better readability.	docs/unit-examples-forward/Beltrami_flow.py
Improved code clarity and organization in the heat equation example.	Removed an unnecessary import statement. Fixed a lambda function to improve code clarity.	docs/unit-examples-forward/heat.ipynb
Adjusted imports for better code organization and to support new functionality.	Added import time to _trainer.py. Removed unused jax.numpy import and added from pinnx import utils in geometry_2d.py.	pinnx/_trainer.py pinnx/geometry/geometry_2d.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time. You can also use
  this command to specify where the summary should be inserted.

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey @chaoming0625 - I've reviewed your changes - here's some feedback:

Overall Comments:

• There's a typo in the parameter names: 'measture' should be 'measure' in both the compile() and train() methods
• The commit message 'update docs' doesn't accurately reflect the changes being made - this PR includes functional changes to the trainer class and code reorganization beyond just documentation updates
• Consider implementing the timing functionality using a profiling decorator or callback system instead of adding it directly to the training loop, which would keep the core code cleaner

Here's what I looked at during the review

• 🟡 General issues: 2 issues found
• 🟢 Security: all looks good
• 🟢 Testing: all looks good
• 🟢 Complexity: all looks good
• 🟢 Documentation: all looks good

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

issue (typo): Parameter name contains a typo: 'measture' should be 'measure'

Suggested implementation:

        measure_train_step_compile_time: bool = False,

        if measure_train_step_compile_time:

You may need to:

1. Update any documentation strings that reference this parameter
2. Update any calls to this method that explicitly name this parameter

[sourcery-ai]

issue (bug_risk): Inconsistent return types could cause bugs - method returns tuple when measuring time but self otherwise

Consider using a consistent return type and providing the timing information through a different mechanism, such as a class attribute or logging.