What is the fastest method to set the diagonals of a matrix to 1?

Question

I am trying to replace the zeros along diagonal of a distance matrix to a list of constants. In this case, a diagonal of 1's.

It appears that the LinearAlgebra package no longer exists so I can't use LinearAlgebra`SetMatrixDiagonal as shown in this answer to a previous question.

Here are some methods to do this and their timings on my machine:

mat = RandomReal[{0, 1}, {10000, 10000}];

Adding the identity matrix would be doing a lot of additions of 0, but it is pretty fast:

RepeatedTiming[mat + IdentityMatrix[10000];]
(*{0.032033125, Null}*)

ReplacePart likely suffers from pattern matching:

RepeatedTiming[ReplacePart[mat , {i_, i_} -> 1];]
(*{14.149746, Null}*)

A loop with Set is faster than adding the identity matrix:

RepeatedTiming[Do[mat[[i, i]] = 1, {i, 1, 10000}]]
(*{0.00310015625`,Null}*)

Perhaps Compile might help here?

setdiag = 
  Compile[
    {{mat, _Real, 2}}, 
    Block[{lmat = mat},
      Do[lmat[[i, i]] = 1, {i, 1, Length[mat]}];
      lmat
    ]
  ]
RepeatedTiming[setdiag[mat];]
(*{0.104903`,Null}*)

But it doesn't.

Maybe ReplacePart without patterns?

RepeatedTiming[
  MapThread[ReplacePart[#1, #2 -> 1] &, {mat, Range[Length[mat]]}];
]
(*{2.434717, Null}*)

Another way of using ReplacePart is better:

RepeatedTiming[
 ReplacePart[mat, 
   Thread[Transpose[{Range[Length[mat]], Range[Length[mat]]}] -> 1]];]
(*{0.01133825`,Null}*)

Can anyone find a better way than the procedural Do?

Answer 1

Edit

A "one-line" way without any C source code:

cf2 = FunctionCompile[
    Function[{Typed[a, "PackedArray"["Real64", 2]], Typed[s, "Integer64"]},
        Module[{carr, len = s*s},
                carr = Array`GetData[a]
                ; Do[ToRawPointer[carr, i, 1.], {i, 0, len - 1, s + 1}]
                ;
            ]]]

RepeatedTiming[cf2[mat, 10000];, 5]
(* {0.0000728393, Null} *)

Things even get better: Using Parallel`ParallelDo instead of Do, we can boost our performance further:

RepeatedTiming[cf3[mat, 10000];, 5]
(* {0.0000263171, Null} *)

---

As OP already suspected, new features of the compiler in 13.1 (like LibraryFunctionDeclaration, RawPointer, etc.) provides an alternative and cleaner way than LibraryLink. The following setup is basically the same as the compilerDemoBase.c example from ToRawPointer's doc page.

1. The C code is as simple as:

#include "WolframLibrary.h"
DLLEXPORT int set_diag_one(double* in, long long s) {
   long len = s*s;
   for (long i = 0; i < len; i += s+1) *(in+i) = 1;
   return 0;
}

Store the code as string in src and compile it:

CreateLibrary[src, "setDiag"]

2. Declare the external function with LibraryFunctionDeclaration:

funcDec = LibraryFunctionDeclaration[
    "set_diag_one", "setDiag",
    {"RawPointer"::["CDouble"], "CLongLong"} -> "CInt"
];

3. Use it in FunctionCompile. Note accroding to Possible Issues in ToRawPointer's doc page, it will cause a value copy. So without trying ToRawPointer, here Array`GetData(*) is used instead to get the raw pointer to a's underline data.

_{* See line 81 in ...\13.1\SystemFiles\Components\Compile\TypeSystem\Declarations\RectangularArray\DenseArray\NumericArray.m.}

cf = FunctionCompile[funcDec,
    Function[
        {Typed[a, "PackedArray"["Real64", 2]], Typed[s, "Integer64"]},
        Module[{ptr},
                ptr = Cast[Array`GetData[a], "RawPointer"::["CDouble"], "BitCast"];
                LibraryFunction["set_diag_one"][ptr, s]
            ]]]

The performance of cf is slightly better than previous solution based on LibraryLink (about 0.00008 s VS 0.0001 s).

Solution described here:

mat = RandomReal[{0, 1}, {10000, 10000}] // Developer`ToPackedArray;
RepeatedTiming[cf[mat, 10000];, 5]
(* {0.000074689, Null} *)
Diagonal[mat] // Union
(* {1.} *)

Solution based on LibraryLink:

RepeatedTiming[setDiag[];, 5]
(* {0.0000996299, Null} *)

Answer 2

A long comment...I'm using V13.1, Mac M1 Pro, which might matter...

First, I would point out the "About" section of the profile of user @WReach. Toward the end, the question "Why don't you like to answer questions about Mathematica performance?" is answered. There are many good points, but I'm often rash enough to answer such questions anyway, when I'm not confused.

Second, while @WReach does not mention CPUs (I take the use of "system" to mean Mma), I know just enough about CPUs to get confused and not enough to get unconfused. I know CPUs use heuristics and caches to accelerate some repetitive actions. When does artificially repeating an operation artificially lower the timing? I have no real understanding of this. Can it affect RepeatedTiming[f[x]] or AbsoluteTiming[Do[f[x], {100}]]? [Odd performance I discovered about my Mac M1: For large $N$, it can add (resp. multiply) $N$ complex numbers about 15% faster than it can add (resp. multiply) $2N$ real numbers; adding and multiplying take the same time in both cases. Bully for Apple, you might say. CPUs are more sophisticated than I understand.]

Third, to get to the OP's examples, setting mat[[1, 1]] to 1 takes much, much longer than setting it to 1.. A packed Real matrix will be unpacked (i.e. copied into a less efficient form) if you set any of its elements to an Integer. RepeatedTiming[Do[...]] hides this because RepeatedTiming returns a trimmed mean, which discards the long time the first Do loop trial takes; in subsequent trials, mat is already unpacked and Do[...] is efficient.

Fourth, on my machine, unpacking mat takes 2 or more seconds; copying a packed mat takes about 0.034s.

mat = RandomReal[{0, 1}, {10000, 10000}];

(foo = mat; foo[[1, 1]] = 1.); // RepeatedTiming
(* copies a packed mat into foo *)
(*  {0.034286, Null}  *)

(foo = mat; foo[[1, 1]] = 1); // RepeatedTiming
(* unpacks mat into foo *)    
(*  {2.01679, Null}  *)

mat // Developer`FromPackedArray; // RepeatedTiming
(* simply unpacks mat -- why less efficient I don't know *)
(*  {2.60058, Null}  *)

Fifth, it would be interesting if an actual (complete) application did not have to copy mat in some form, given that its values are being changed. It's conceivable, and given its size, it's desirable.

Sixth, Does Do[mat[[i, i]] = 1., {i, 1, Length@mat}] do an in-place update of mat without copying it? It seems that it should, but on my machine, it seems not if a new mat is used in each trial:

Table[
  mat = RandomReal[{0, 1}, {10000, 10000}]; 
  First@AbsoluteTiming[Do[mat[[i, i]] = 1., {i, 1, 10000}]],
  {100}] // Mean
(*  0.0357951  <-- roughly = time to copy mat *)

mat = RandomReal[{0, 1}, {10000, 10000}];
Table[
  First@AbsoluteTiming[Do[mat[[i, i]] = 1., {i, 1, 10000}]],
  {100}] // Mean
(*  0.00406767  *)

Seventh, using 1. instead of 1 changes some of the ReplacePart timings as well. Here's a variant on the OP's that takes just a bit longer than it does to copy mat:

mat = RandomReal[{0, 1}, {10000, 10000}];
RepeatedTiming[ReplacePart[mat, Table[{i, i} -> 1., {i, Length@mat}]];]
(*  {0.0498069, Null}  *)

Table[
  mat = RandomReal[{0, 1}, {10000, 10000}]; 
  First@AbsoluteTiming[
    ReplacePart[mat, Table[{i, i} -> 1., {i, Length@mat}]]],
  {30}] // Mean
(*  0.0383822  *)

Eighth, the LinearAlgebra`Private`SetMatrixDiagonal function takes about 0.048s, or a little slower than Do[]. Interestingly, it does not matter whether the diagonal is set to ConstantArray[1, Length[mat]] or to ConstantArray[1., Length[mat]]. Even with the integer 1, the diagonal is set to the real 1..

Maybe an answer: My guess is that the in-place Do[...] is fastest and that when mat is copied. The fastest timing of 0.004s can be achieved if mat has been "read" into memory, as done by mat + 1.; below. Of course, that read operation takes time, so it's not that you can get the timing down to 0.004s for free. Note if the line is changed to mat = mat + 1.; in both places, then the difference in timing increases to 0.049s. If other operations come between the "read" of mat and the Do loop, then the timing can increase.

Mean@Table[
   mat = RandomReal[{0, 1}, {10000, 10000}];
   First@AbsoluteTiming[
     mat + 1.;
     Do[mat[[i, i]] = 1., {i, 1, 10000}]],
   {30}] -
 Mean@Table[
   mat = RandomReal[{0, 1}, {10000, 10000}];
   First@AbsoluteTiming[mat + 1.],
   {30}]
(*  0.00451733  *)

Summary: Probably @WReach is right about the difficulty of nailing down precisely what is the best performance. The biggest issue here is 1 vs. 1., that is, unpacking vs. keeping mat packed. The next biggest is probably whether mat gets copied. That seems tricky. Maybe someone else has a clear view on when.

Answer 3

It is pretty fast to use LibraryLink, making sure the memory is shared between the library and the Mathematica kernel. Note that you need a C compiler. C code:

/*
 Based on demo_shared.c, the example about shared memory management for 
 communicating between Mathematica and a DLL.
*/

#include "WolframLibrary.h"

static MTensor tensor;

DLLEXPORT mint WolframLibrary_getVersion( ) {
    return WolframLibraryVersion;
}

DLLEXPORT int WolframLibrary_initialize(WolframLibraryData libData) {
    return 0;
}

DLLEXPORT int loadArray(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {
    tensor = MArgument_getMTensor(Args[0]);
    MArgument_setInteger(Res, 0);
    return 0;
}

DLLEXPORT int setElement(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {
    mint pos[2];
    mint len = MArgument_getInteger(Args[0]);
    
    for(int ii = 1; ii <= len; ii ++){
        pos[0] = ii; pos[1] = ii;
        libData->MTensor_setReal(tensor, pos, 1);
    }
    
    MArgument_setInteger(Res, 0);
    return 0;
}

DLLEXPORT int unloadArray(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {
    libData->MTensor_disown(tensor);
    MArgument_setInteger(Res, 0);
    return 0;
}

Mathematica code (note that you need to set pathToCodeFile):

<<CCompilerDriver`
mat = RandomReal[{0, 1}, {10000, 10000}];
setDiagLib = CreateLibrary[{pathToCodeFile}, "setDiagLib"];
loadFun = 
  LibraryFunctionLoad[setDiagLib, 
   "loadArray", {{Real, _, "Shared"}}, Integer];
unloadFun =
  LibraryFunctionLoad[setDiagLib, "unloadArray", {}, Integer];
setFun = 
  LibraryFunctionLoad[setDiagLib, "setElement", {Integer}, 
   Integer];
setDiag[]:=
  (
     loadFun[mat];
     setFun[Length@mat];
     unloadFun[];);

Timing

setDiag[]//RepeatedTiming//First
mat[[1,1]]

0.000159
1.