Compute randomized Kronecker product and dump it on the fly.

  A standard Kronecker product between matrices A and B produces
                        [[a_11 B, ..., a_1n B],
                                  ...
                         [a_m1 B, ..., a_mn B]]
    (if A's size is (m, n) and B's size is (p, q) then A Kronecker B has size
    (m p, n q)).
    Here we modify the standard Kronecker product expanding matrices in
    https://cs.stanford.edu/~jure/pubs/kronecker-jmlr10.pdf
    and randomize each block-wise operation a_ij B in the Kronecker product as
    in https://arxiv.org/pdf/1901.08910.pdf section III.4.
    The matrix we produce is
                       [[F(a_11, B, w_11), ..., F(a_1n, B, w_1n)],
                                           ...
                       [F(a_m1, B, w_m1), ... , F(a_mn, B, w_mn)]]
    where (w_ij) is a sequence of pseudo random numbers and F is randomized
    operator which will:
      1) Shuffle rows and columns of B independently at random;
      2) Dropout elements of B with a rate 1 - a_ij to compute
        F(a_ij, B, w_ij).
    (It is noteworthy that there is an abuse of notation above when writing
    F(a_ij, B, w_ij) as each block-wise operation will in fact consume
    multiple elements of the sequence (w_ij)).

  Args:
    left_matrix: sparse SciPy csr matrix with values in [0, 1].
    right_matrix: sparse SciPy coo signed binary matrix. +1 values correspond
      to train set and -1 values correspond to test set.
    train_indices_out_path: path to output train file. The non zero indices of
      the resulting sparse matrix are dumped as a series of pickled records.
      As many shard will be created as there are rows in left matrix. The shard
      corresponding to row i in the left matrix has the suffix _i appended to
      its file name. Each shard contains a pickled list of list each of which
      corresponds to a users.
    test_indices_out_path: path to output train file. The non zero indices of
      the resulting sparse matrix are dumped as a series of pickled records.
      As many shard will be created as there are rows in left matrix. The shard
      corresponding to row i in the left matrix has the suffix _i appended to
      its file name. Each shard contains a pickled list of list each of which
      corresponds to a users.
    train_metadata_out_path: path to optional complementary output file
      containing the number of train rows (r), columns (c) and non zeros (nnz)
      in a pickled SparseMatrixMetadata named tuple.
    test_metadata_out_path: path to optional complementary output file
      containing the number of test rows (r), columns (c) and non zeros (nnz)
      in a pickled SparseMatrixMetadata named tuple.
    remove_empty_rows: whether to remove rows from the synthetic train and
      test matrices which are not present in the train or the test matrix.

  Returns:
    (metadata, train_metadata, test_metadata) triplet of SparseMatrixMetadata
      corresponding to the overall data set, train data set and test data set.

def output_randomized_kronecker_to_pickle(
    left_matrix, right_matrix,
    train_indices_out_path, test_indices_out_path,
    train_metadata_out_path=None, test_metadata_out_path=None,
    remove_empty_rows=True):
  """Compute randomized Kronecker product and dump it on the fly.

  A standard Kronecker product between matrices A and B produces
                        [[a_11 B, ..., a_1n B],
                                  ...
                         [a_m1 B, ..., a_mn B]]
    (if A's size is (m, n) and B's size is (p, q) then A Kronecker B has size
    (m p, n q)).
    Here we modify the standard Kronecker product expanding matrices in
    https://cs.stanford.edu/~jure/pubs/kronecker-jmlr10.pdf
    and randomize each block-wise operation a_ij B in the Kronecker product as
    in https://arxiv.org/pdf/1901.08910.pdf section III.4.
    The matrix we produce is
                       [[F(a_11, B, w_11), ..., F(a_1n, B, w_1n)],
                                           ...
                       [F(a_m1, B, w_m1), ... , F(a_mn, B, w_mn)]]
    where (w_ij) is a sequence of pseudo random numbers and F is randomized
    operator which will:
      1) Shuffle rows and columns of B independently at random;
      2) Dropout elements of B with a rate 1 - a_ij to compute
        F(a_ij, B, w_ij).
    (It is noteworthy that there is an abuse of notation above when writing
    F(a_ij, B, w_ij) as each block-wise operation will in fact consume
    multiple elements of the sequence (w_ij)).

  Args:
    left_matrix: sparse SciPy csr matrix with values in [0, 1].
    right_matrix: sparse SciPy coo signed binary matrix. +1 values correspond
      to train set and -1 values correspond to test set.
    train_indices_out_path: path to output train file. The non zero indices of
      the resulting sparse matrix are dumped as a series of pickled records.
      As many shard will be created as there are rows in left matrix. The shard
      corresponding to row i in the left matrix has the suffix _i appended to
      its file name. Each shard contains a pickled list of list each of which
      corresponds to a users.
    test_indices_out_path: path to output train file. The non zero indices of
      the resulting sparse matrix are dumped as a series of pickled records.
      As many shard will be created as there are rows in left matrix. The shard
      corresponding to row i in the left matrix has the suffix _i appended to
      its file name. Each shard contains a pickled list of list each of which
      corresponds to a users.
    train_metadata_out_path: path to optional complementary output file
      containing the number of train rows (r), columns (c) and non zeros (nnz)
      in a pickled SparseMatrixMetadata named tuple.
    test_metadata_out_path: path to optional complementary output file
      containing the number of test rows (r), columns (c) and non zeros (nnz)
      in a pickled SparseMatrixMetadata named tuple.
    remove_empty_rows: whether to remove rows from the synthetic train and
      test matrices which are not present in the train or the test matrix.

  Returns:
    (metadata, train_metadata, test_metadata) triplet of SparseMatrixMetadata
      corresponding to the overall data set, train data set and test data set.
  """
  logging.info("Writing item sequences to pickle files %s and %s.",
               train_indices_out_path, test_indices_out_path)

  num_rows = 0
  num_removed_rows = 0
  num_cols = left_matrix.shape[1] * right_matrix.shape[1]
  num_interactions = 0

  num_train_interactions = 0
  num_test_interactions = 0

  if not set(right_matrix.data).issubset({-1, 1}):
    raise ValueError(
        "Values of sparse matrix should be -1 or 1 but are:",
        set(right_matrix.data))

  for i in xrange(left_matrix.shape[0]):

    (shard_num_removed_rows, shard_metadata, shard_train_metadata,
     shard_test_metadata) = _compute_and_write_row_block(
         i, left_matrix, right_matrix, train_indices_out_path,
         test_indices_out_path, remove_empty_rows)

    num_rows += shard_metadata.num_rows
    num_removed_rows += shard_num_removed_rows
    num_interactions += shard_metadata.num_interactions
    num_train_interactions += shard_train_metadata.num_interactions
    num_test_interactions += shard_test_metadata.num_interactions

    logging.info("%d total interactions written.", num_interactions)
    logging.info("%d total rows removed.", num_removed_rows)
    logging.info("%d total train interactions written.", num_train_interactions)
    logging.info("%d toal test interactions written.", num_test_interactions)

  logging.info("Done writing.")

  metadata = SparseMatrixMetadata(
      num_interactions=num_interactions,
      num_rows=num_rows, num_cols=num_cols)
  train_metadata = SparseMatrixMetadata(
      num_interactions=num_train_interactions,
      num_rows=num_rows, num_cols=num_cols)
  test_metadata = SparseMatrixMetadata(
      num_interactions=num_test_interactions,
      num_rows=num_rows, num_cols=num_cols)

  if train_metadata_out_path is not None:
    util.write_metadata_to_file(
        train_metadata, train_metadata_out_path, tag="train")
  if test_metadata_out_path is not None:
    util.write_metadata_to_file(
        test_metadata, test_metadata_out_path, tag="test")

  return metadata, train_metadata, test_metadata