Files
wt-tut-presentations/src/2026-01-30/gen_correlated_data.py

39 lines
943 B
Python

import numpy as np
import matplotlib.pyplot as plt
from numpy.typing import NDArray
import argparse
def twodim_array_to_pgfplots_table_string(a: NDArray):
return (
" \\\\\n".join([" ".join([str(vali) for vali in val]) for val in a]) + "\\\\\n"
)
def main():
# Parse command line arguments
parser = argparse.ArgumentParser()
parser.add_argument("--correlation", "-c", type=np.float32, required=True)
parser.add_argument("-N", type=np.int32, required=True)
parser.add_argument("--plot", "-p", action="store_true")
args = parser.parse_args()
# Generate & plot data
means = np.array([0, 0])
cov = np.array([[1, args.correlation], [args.correlation, 1]])
x = np.random.multivariate_normal(means, cov, size=args.N)
print(twodim_array_to_pgfplots_table_string(x))
if args.plot:
plt.scatter(x[:, 0], x[:, 1])
plt.show()
if __name__ == "__main__":
main()