Add zip() for vectorizing matrix lookups for paired indices

[stekien]

Summary

Adds stan::math::zip(x, idx_row, idx_col), which returns the column vector whose i-th element is x(idx_row[i], idx_col[i]) for 1-based index arrays. This is the paired-index ("zipped") gather that currently has to be written as an explicit loop in Stan programs when a model indexes a matrix by two parallel index arrays — e.g. a J x K parameter matrix indexed by aligned location[n] / time[n] arrays.

The implementation converts the two index arrays into a single linear index via Eigen::Map over the input std::vectors and applies it to x.reshaped(), so the gather is one vectorized indexed access rather than an element-wise loop. Argument checking uses the existing check_size_match and check_range utilities: the two index arrays must have equal length, and check_range is applied to minCoeff()/maxCoeff() of each index array so the bounds check costs two passes rather than one check per element.

The only change to an existing file is registering the new header in stan/math/prim/fun.hpp.

Related to issue stan-dev/stanc3#1704

Tests

New file test/unit/math/prim/fun/zip_test.cpp with two tests:

MathFunctions.zip — expected behavior. Builds a 5x4 Eigen::MatrixXd whose entries encode their own position (delta(i,j) == 10*(i+1) + (j+1)) and zips nine index pairs that include repeated pairs, out-of-order access, and both boundary rows and columns. Asserts the result size equals the index array length and checks every element against the corresponding delta(i, j).
MathFunctions.zip_errors — error handling. Asserts valid input does not throw, then covers each rejection path: mismatched index array lengths throws std::invalid_argument; row index above x.rows(), row index 0, column index above x.cols(), and negative column index each throw std::out_of_range.
Run with ./runTests.py test/unit/math/prim/fun/zip_test.cpp.

Side Effects

None for existing code. The change is one new header plus one #include line in stan/math/prim/fun.hpp; no existing function signatures or behavior change. The added header is small and pulls in only meta.hpp, err.hpp, fun/Eigen.hpp and , all already transitively included by prim.hpp, so the compile-time cost is negligible.

Scope limitations worth flagging for review:

The overload is constrained to require_eigen_matrix_dynamic_t, so it accepts any Eigen dynamic matrix including Eigen::Matrix<var, -1, -1>, but var_valueEigen::MatrixXd (varmat) is not an Eigen type and is not supported. Only prim/double tests are included; no rev or fwd tests yet.
Indices are std::vector only, matching how Stan passes array[] int.

Release notes

Added zip(matrix, array[] int, array[] int), which returns the vector of matrix elements selected by a pair of aligned 1-based row and column index arrays.

Checklist

• Copyright holder: Simons Foundation

By submitting this pull request, the copyright holder is agreeing to the license the submitted work under the following licenses:
- Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
- Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

• the basic tests are passing

  • unit tests pass (to run, use: ./runTests.py test/unit)
  • header checks pass, (make test-headers)
  • dependencies checks pass, (make test-math-dependencies)
  • docs build, (make doxygen)
  • code passes the built in C++ standards checks (make cpplint)

• the code is written in idiomatic C++ and changes are documented in the doxygen

• the new changes are tested

[stan-buildbot]

Name	Old Result	New Result	Ratio	Performance change( 1 - new / old )
stat_comp_benchmarks/benchmarks/gp_regr/gp_regr.stan	0.23	0.22	1.01	0.54% faster
stat_comp_benchmarks/benchmarks/gp_regr/gen_gp_data.stan	0.06	0.06	0.98	-2.26% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix/low_dim_gauss_mix.stan	6.31	6.32	1.0	-0.07% slower
stat_comp_benchmarks/benchmarks/low_dim_corr_gauss/low_dim_corr_gauss.stan	0.02	0.02	1.03	2.91% faster
stat_comp_benchmarks/benchmarks/irt_2pl/irt_2pl.stan	8.46	8.44	1.0	0.27% faster
stat_comp_benchmarks/benchmarks/gp_pois_regr/gp_pois_regr.stan	4.53	4.49	1.01	0.74% faster
stat_comp_benchmarks/benchmarks/sir/sir.stan	166.72	165.71	1.01	0.61% faster
stat_comp_benchmarks/benchmarks/garch/garch.stan	0.92	0.89	1.03	3.14% faster
stat_comp_benchmarks/benchmarks/arma/arma.stan	0.7	0.71	0.99	-0.98% slower
stat_comp_benchmarks/benchmarks/pkpd/one_comp_mm_elim_abs.stan	42.86	42.56	1.01	0.69% faster
stat_comp_benchmarks/benchmarks/pkpd/sim_one_comp_mm_elim_abs.stan	0.59	0.6	1.0	-0.18% slower
stat_comp_benchmarks/benchmarks/low_dim_gauss_mix_collapse/low_dim_gauss_mix_collapse.stan	20.82	20.78	1.0	0.19% faster
stat_comp_benchmarks/benchmarks/eight_schools/eight_schools.stan	0.11	0.11	1.02	1.77% faster
stat_comp_benchmarks/benchmarks/arK/arK.stan	3.19	3.22	0.99	-1.11% slower
performance.compilation	376.7	377.86	1.0	-0.31% slower
Mean result: 1.0041787509034354

Jenkins Console Log
Jenkins Build Stages
Commit hash: df9c251315a66d40077fa83739277fd6c77e2cd4

Machine information

Distributor ID:	Ubuntu
Description:	Ubuntu 20.04.3 LTS
Release:	20.04
Codename:	focal

CPU:

Architecture:                            x86_64
CPU op-mode(s):                          32-bit, 64-bit
Byte Order:                              Little Endian
Address sizes:                           43 bits physical, 48 bits virtual
CPU(s):                                  256
On-line CPU(s) list:                     0-255
Thread(s) per core:                      2
Core(s) per socket:                      64
Socket(s):                               2
NUMA node(s):                            2
Vendor ID:                               AuthenticAMD
CPU family:                              23
Model:                                   49
Model name:                              AMD EPYC 7742 64-Core Processor
Stepping:                                0
Frequency boost:                         enabled
CPU MHz:                                 1495.355
CPU max MHz:                             3416.0681
CPU min MHz:                             1500.0000
BogoMIPS:                                4491.85
Virtualization:                          AMD-V
L1d cache:                               4 MiB
L1i cache:                               4 MiB
L2 cache:                                64 MiB
L3 cache:                                512 MiB
NUMA node0 CPU(s):                       0-63,128-191
NUMA node1 CPU(s):                       64-127,192-255
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Old microcode:             Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec rstack overflow:      Mitigation; Safe RET
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; Retpolines; IBPB conditional; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Mitigation; IBPB before exit to userspace
Flags:                                   fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sev sev_es

G++:

g++ (Ubuntu 9.4.0-1ubuntu1~20.04) 9.4.0
Copyright (C) 2019 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

Clang:

clang version 10.0.0-4ubuntu1 
Target: x86_64-pc-linux-gnu
Thread model: posix
InstalledDir: /usr/bin

[WardBrian]

I think this was a misreading of the request in stan-dev/stanc3#1704 and should be reverted. The zip function should just construct an array of pairs, and all the current indexing code lives in stan-dev/stan not stan-dev/math.