feat: add memory profiling and JSON output to scaling benchmark

[yuvraajnarula]

Describe your changes

This PR extends the existing scaling benchmark (graph_creation_scaling.py) with memory profiling, machine‑readable output (JSON), and a separate memory scaling plot. It builds on the runtime scaling work in #117 and addresses the first part of #137 (memory profiling and baseline generation).

Key additions:

• --track-memory - enables peak memory tracking via tracemalloc.
• --output-json <file> - saves raw results (grid points, runtime, peak memory) as JSON for later use (e.g., CI regression tests).
• --output-plot-memory <file> - generates a dedicated plot of peak memory vs number of grid points (requires --track-memory).
• Refactored the script into reusable functions:
  • run_benchmark() – returns a list of dicts with the measured data.
  • plot_runtime_scaling() - creates the runtime vs grid points plot.
  • plot_memory_scaling() - creates the memory vs grid points plot.

Motivation and context:

• Establishing baseline memory/runtime numbers on main is required for future performance regression testing.
• Machine‑readable output (JSON) enables programmatic consumption, which will be used in a follow‑up PR for CI regression tests.
• The memory scaling plot helps identify when graph generation becomes memory‑bound, critical for large global grids (e.g., 10⁷ nodes).

Dependencies:

• tracemalloc (standard library) - no new external dependencies.

Issue Link

Closes #137 (first part - memory profiling and baseline tooling)

Type of change

• 🐛 Bug fix
• ✨ New feature
• 💥 Breaking change
• 📖 Documentation

Checklist before requesting a review

• My branch is up-to-date with the target branch (mllam/main).
• I have performed a self-review of my code.
• For any new/modified functions/classes I have added docstrings.
• I have placed in-line comments to clarify the intent.
• I have updated the documentation (README.md and docstrings).
• I have added tests that prove my fix is effective or that my feature works (existing tests pass; new functionality tested manually).
• I have given the PR a name that clearly describes the change, written in imperative form: "feat: add memory profiling and JSON output to scaling benchmark".
• I have requested a reviewer and an assignee (assignee is responsible for merging).

Checklist for reviewers

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in the [Unreleased] section under added.

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[leifdenby]

Thank you for writing this @yuvraajnarula ! And sorry for taking so long to look at it.

Would be able to post a few terminal printouts or screenshots of you running this so I can see how it is supposed to work?

[prajwal-tech07]

Reviewed and tested the output format - looks good. --output-json produces a clean list of {grid_points, runtime_s, peak_memory_mb} records, which is exactly what the CI regression check (#144) will consume in compare.py. Verified the archetype CLI choices map to real functions (create_keisler_graph / create_graphcast_graph / create_oskarsson_hierarchical_graph) and that create_fn(coords=xy) matches their signatures. 👍

Two small non-blocking suggestions:

• The README still documents the old CLI (--output → scaling_plot.png, only --show). Could it be updated to the current flags - especially the new --output-json, --track-memory, and --output-plot-memory, since the machine-readable output is the main point here?
• The description says "Closes #137" - since this is only the first part, maybe "Part of #137" so the regression-testing tracking stays open (continued in #144).

Nice work @yuvraajnarula. 🚀

[yuvraajnarula]

Thanks for the review and for testing the output format, @prajwal-tech07!

Good catch on the README and the description. I'll get the CLI flags updated in the docs and change the copy to 'Part of #137' so we keep the automated tracking open for the CI work.

Saw your sequencing question over on #144 as well, and happy to wait for @leifdenby to chime in on the merge order, or adjust/rebase this branch however is easiest.

[leifdenby]

@yuvraajnarula I've merged in #117. Do you want to update your PR branch with main and @prajwal-tech07 and I can do a last review so we can merge this in? Also, could you add a changelog entry? :) THanks!

[leifdenby]

@yuvraajnarula because @prajwal-tech07 wants to build on your nice work here I will resolve the merge conflicts and then we will get this merged in today. Hope that is ok, and thank you again for your contribution!

[prajwal-tech07]

Thanks for resolving the conflicts, Leif - the merge looks correct to me. Nothing was lost from #117 (all the --min-N/--max-N/--num-steps/--archetype args, the choices= constraint, the create_fake_xy harness and the O(N) reference line are all still there), no leftover conflict markers, and the changelog reads well. Happy for this to go in today.

One non-blocking thing I noticed for a quick follow-up (doesn't need to hold up the merge since it's a hand-run benchmark script): the new memory-plot output doesn't actually write a file yet. output_path is passed into both plot_runtime_scaling and plot_memory_scaling but never used, and the only savefig in main targets args.output_plot_runtime. So --output-plot-memory produces no file, and because plot_memory_scaling opens its own figure, running with --track-memory --output-plot-memory saves the memory plot into the runtime PNG.

Fix is ~5 lines - add plt.savefig(output_path) at the end of each plot function and drop the lone plt.savefig(...) in main (plus a ylabel/legend on the runtime plot). Since it's tiny, feel free to squeeze it into this PR before merging if easy; otherwise I'm building on this benchmark anyway and happy to send it as a small follow-up right after. Either way, no objection to merging today.

[leifdenby]

One non-blocking thing I noticed for a quick follow-up (doesn't need to hold up the merge since it's a hand-run benchmark script): the new memory-plot output doesn't actually write a file yet. output_path is passed into both plot_runtime_scaling and plot_memory_scaling but never used, and the only savefig in main targets args.output_plot_runtime. So --output-plot-memory produces no file, and because plot_memory_scaling opens its own figure, running with --track-memory --output-plot-memory saves the memory plot into the runtime PNG.

Ah, good catch! I will fix that. Thanks

[prajwal-tech07]

Fix looks perfect - each plot function now saves to its own output_path and the stray savefig in main is gone, so --output-plot-memory writes its file and no longer clobbers the runtime plot. Verified nothing else regressed and the conflict resolution is still intact. LGTM, good to merge 🚀