title: Multiplying large matrices with dask
body: <p>I am working on a project which basically boils down to solving the matrix equation</p>
  
  <pre><code>A.dot(x) = d
  </code></pre>
  
  <p>where <code>A</code> is a matrix with dimensions roughly 10 000 000 by 2000 (I would like to increase this in both directions eventually). </p>
  
  <p><code>A</code> obviously does not fit in memory, so this has to be parallelized. I do that by solving <code>A.T.dot(A).dot(x) = A.T.dot(d)</code> instead. <code>A.T</code> will have dimensions 2000 by 2000. It can be calculated by dividing <code>A</code> and <code>d</code> into chunks <code>A_i</code> and <code>d_i</code>, along the rows, calculate <code>A_i.T.dot(A_i)</code> and <code>A_i.T.dot(d_i)</code>, and sum these. Perfect for parallellizing. I have been able to implement this with the multiprocessing module, but it is 1) hard to scale any further (increasing <code>A</code> in both dimensions), due to memory use, and 2) not pretty (and therefore not easy to maintain).</p>
  
  <p>Dask seems to be a very promising library for solving both these problems, and I have made some attempts. My <code>A</code> matrix is complicated to calculate: It is based on approximately 15 different arrays (with size equal to the number of rows in <code>A</code>), and some are used in an iterative algorithm to evaluate associated Legendre function. When the chunks are small (10000 rows), it takes a very long time to build the task graph, and it takes a lot of memory (the memory increase coincides with the call to the iterative algorithm). When the chunks are larger (50000 rows), the memory consumption prior to calculations is a lot smaller, but it is rapidly depleted when calculating <code>A.T.dot(A)</code>. I have tried with <code>cache.Chest</code>, but it significantly slows down the calculations. </p>
  
  <p>The task graph must be very large and complicated - calling <code>A._visualize()</code> crashes. With simpler <code>A</code> matrices, it works to do this directly (see response by @MRocklin). Is there a way for me to simplify it?</p>
  
  <p>Any advice on how to get around this would be highly appreciated.</p>
  
  <p><strike>As a toy example, I tried </p>
  
  <pre><code>A = da.ones((2e3, 1e7), chunks = (2e3, 1e3)) 
  (A.T.dot(A)).compute()
  </code></pre>
  
  <p>This also failed, using up all the memory with only one core being active. With <code>chunks = (2e3, 1e5)</code>, all cores start almost immediately, but <code>MemoryError</code> appears within 1 second (I have 15 GB on my current computer). <code>chunks = (2e3, 1e4)</code> was more promising, but it ended up consuming all memory as well. </strike></p>
  
  <p>Edit:
  I struckthrough the toy example test, because the dimensions were wrong, and corrected the dimensions in the rest. As @MRocklin says, it does work with the right dimensions. I added a question which I now think is more relevant to my problem.</p>
  
  <p>Edit2:
  This is a much simplified example of what I was trying to do. The problem is, I believe, the recursions involved in defining the columns in <code>A</code>. </p>
  
  <pre><code>import dask.array as da
  
  N = 1e6
  M = 500
  
  x = da.random.random((N, 1), chunks = 5*M)
  
  # my actual 
  A_dict = {0:x}
  for i in range(1, M):
      A_dict[i] = 2*A_dict[i-1]
  A = da.hstack(tuple(A_dict.values()))
  A = A.rechunk((M*5, M))
  ATA = A.T.dot(A)
  </code></pre>
  
  <p>This seems to lead to a very complicated task graph, which takes up a lot of memory before the calculations even start.</p>
  
  <p>I have now solved this by placing the recursion in a function, with <code>numpy</code> arrays, and more or less do <code>A = x.map_blocks(...)</code>. </p>
  
  <p>As a second note, once I have the <code>A</code> matrix task graph, calculating <code>A.T.dot(A)</code> directly does seem to give some memory issues (memory usage is not very stable). I therefore explicitly calculate it in chunks, and sum the results. Even with these workarounds, dask makes a big difference in speed and readability.</p>