updates to incorporate the matrix chunking method #79
Conversation
There was a problem hiding this comment.
Thanks @ccain002! This is promising. I've made a few comments to get the code into shape, then we can do some proper unit tests.
src/matvis/cli.py
Outdated
| def get_matrix_sets(bls, ndecimals: int = 2): | ||
| """Find redundant baselines.""" | ||
| uvbins = set() | ||
| msets = [] | ||
|
|
||
| # Everything here is in wavelengths | ||
| bls = np.round(bls, decimals=ndecimals) | ||
| nant = bls.shape[0] | ||
|
|
||
| # group redundant baselines | ||
| for i in range(nant): | ||
| for j in range(i + 1, nant): | ||
| u, v = bls[i, j] | ||
| if (u, v) not in uvbins and (-u, -v) not in uvbins: | ||
| uvbins.add((u, v)) | ||
| msets.append([np.array([i]), np.array([j])) | ||
|
|
||
| return msets |
There was a problem hiding this comment.
I think we can replace this with the little algorithm I had in my other branch. Also, put this in a new module (like "submatrices.py" or something)
src/matvis/cli.py
Outdated
| @@ -290,8 +315,10 @@ def hera_profile(hex_num, nside, keep_ants, outriggers, **kwargs): | |||
|
|
|||
| bls = antpos[np.newaxis, :, :2] - antpos[:, np.newaxis, :2] | |||
| pairs = np.array(get_redundancies(bls.value)) | |||
| matsets = list(get_matrix_sets(bls.value)) | |||
There was a problem hiding this comment.
here, you could accept a string value for matsets that specifies a function for how to divide up the sets
There was a problem hiding this comment.
good idea. This section of the code can expand to include different functions that we populate into submatrices.py.
tests/test_matprod.py
Outdated
| import sys | ||
|
|
||
| sys.path.insert( | ||
| 0, | ||
| "/home/lab-admin/Desktop/Desktop/Graduate_School/Research/ASU/21_group/matvis/src/", | ||
| ) | ||
|
|
There was a problem hiding this comment.
This needs to be removed. You might want to install at the top level with pip install -e . so that when running pytest it can find the library.
There was a problem hiding this comment.
My bad on the mess with this file. I had meant to replace this with a cleaned up version and evidently forgot to.
tests/test_matprod.py
Outdated
| if matsets and method.startswith("CPU"): | ||
| matsets = [ | ||
| (np.array([0, 1, 2, 3]), np.array([0, 1, 2, 3])), | ||
| (np.array([0, 1, 2, 3]), np.array([3, 4])), | ||
| (np.array([3, 4]), np.array([0, 1, 2, 3])), | ||
| (np.array([3, 4]), np.array([3, 4])), | ||
| ] | ||
| elif matsets and method.startswith("GPU"): | ||
| matsets = [ | ||
| (cp.array([0, 1, 2, 3]), cp.array([0, 1, 2, 3])), | ||
| (cp.array([0, 1, 2, 3]), cp.array([3, 4])), | ||
| (cp.array([3, 4]), cp.array([0, 1, 2, 3])), | ||
| (cp.array([3, 4]), cp.array([3, 4])), |
There was a problem hiding this comment.
instead of making them cupy arrays here, convert them in the matprod class respectively.
There was a problem hiding this comment.
I think I figured out how to do this, but I had to use the HAVE_GPU flag instead of use_gpu to specify this, which would cause a problem if the code detected a GPU but was running in CPU mode (I don't know if this would happen).
| for _, (ai, aj) in enumerate(self.antpairs): | ||
| antpairs_set.add((ai, aj)) |
There was a problem hiding this comment.
I think you could simply do this as antpairs_set = set(self.antpairs)
This update adds the capability to use the matrix chunking method to do matrix multiplication, alongside the full matrix product and vector-vector method. This mode takes a new argument "matsets", which is a list of tuples of arrays of rows and columns, which specify sub-matrices. After the sub-matrices are multiplied, the unique pairs are selected from the results (as specified by antpairs) and the output is only the unique pairs, just like in the vector-vector case. I've tested that the new matrix multiplication method itself works properly with the new inputs, but I haven't checked that the updates to the API are correct yet.