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

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 LinearAlgebraSetMatrixDiagonal 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;]
(*{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,
(*{0.01133825,Null}*)


Can anyone find a better way than the procedural Do?

• "It appears that the LinearAlgebra package no longer exists and I can't use LinearAlgebraSetMatrixDiagonal" - it's now LinearAlgebraPrivateSetMatrixDiagonal[] Jul 4 at 14:46
• I am skeptical of your approach involving IdentityMatrix: surely, it is only adding 1 to each element of the diagonal, and not replacing those elements with 1? Jul 4 at 15:49
• The Do approach is fastest because it doesn't involve making a copy. That's going to be hard to beat.
– Jens
Jul 4 at 16:04
• @CraigCarter Yes it's also adding zeros. My point is that, if a diagonal element starts out e.g. as 0.4, then after adding IdentityMatrix, its value will be 1.4, and not 1 as you wanted. Jul 4 at 18:43
• I think a key point here is whether you want a copy of the original matrix with the diagonals replaced (a 'new' matrix) or the original matrix with the diagonal replaced ('in place' modification). As Leonid Shifrin has commented here "The big problem with ReplacePart is that it copies entire list rather than modify a part in-place". Jul 5 at 6:50

## 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 = ArrayGetData[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 ParallelParallelDo 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"]

1. Declare the external function with LibraryFunctionDeclaration:
funcDec = LibraryFunctionDeclaration[
"set_diag_one", "setDiag",
{"RawPointer"::["CDouble"], "CLongLong"} -> "CInt"
];

1. 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 ArrayGetData(*) 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[ArrayGetData[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}] // DeveloperToPackedArray;
RepeatedTiming[cf[mat, 10000];, 5]
(* {0.000074689, Null} *)
Diagonal[mat] // Union
(* {1.} *)


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

• This is great @Silvia. Useful working examples are such a huge help. Thanks. Jul 6 at 18:32
• @CraigCarter You're welcome! Despite the lack of documentation, the new compiler shows some promising potential. I find diving into the source code worth the effort. Jul 6 at 20:01
• Wow, this seems to be the first post diving into the new compiler in this site. Of course +1. Jul 7 at 5:05
• @xzczd Thanks! This was inspired by the recent topic in the official wechat group. :P Jul 7 at 6:43

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 // DeveloperFromPackedArray; // 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 LinearAlgebraPrivateSetMatrixDiagonal 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.

• Thanks. Well, that was quite an education. And, I thought I was asking a simple question. I'm very appreciative of the thoughtful responses. Jul 5 at 8:06

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);
MArgument_setInteger(Res, 0);
return 0;
}

DLLEXPORT int setElement(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {
mint pos;
mint len = MArgument_getInteger(Args);

for(int ii = 1; ii <= len; ii ++){
pos = ii; pos = 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"];
setFun =
Integer];
setDiag[]:=
(
setFun[Length@mat];


Timing

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


0.000159
1.

• This shows a difference: Table[ mat = RandomReal[{0, 1}, {10000, 10000}]; First@AbsoluteTiming[setDiag[];], {30}] // Mean versus Table[ First@AbsoluteTiming[setDiag[];], {30}] // Mean. The first is about as fast (or a little faster) as a primed Do[], about 0.0035s. (The second is about 0.00006s.) Maybe we should write all our Mathematica code in C. :) (+1) Jul 4 at 23:38
• @Jacob Akkerboom, this a useful example. I am guessing the new ToRawPointer in 13.1 may be a way to avoid LibraryFunctionLoad. I wish there were more working examples in the the new compiler features documentation. Jul 5 at 8:19
• @MichaelE2 the reason why the timing is slower when mat is defined inside the loop may be that Mathematica is doing garbage collection for previous discarded data of mat. Perhaps unsurprisingly, adding the following delay seems to be sufficient to prime the pump so to speak.Table[mat = RandomReal[{0, 1}, {10000, 10000}];Pause[0.1];Do[mat[[1,1]] + mat[[1,2]], {ii, 1, 100}];Pause[0.1];First@AbsoluteTiming[setDiag[];], {30}] // Mean. Jul 5 at 19:55
• +1. BTW I'm on Windows with mingw, MTensor_setReal( tensor, &pos, 1) compiles with warning message here, but MTensor_setReal( tensor, pos, 1) works fine. Jul 6 at 12:38
• @Silvia thanks for pointing that out, that was a mistake. Also I guess my solution is not fastest anymore :P. I thought maybe using libData->MTensor_setReal rather than directly setting the data slowed my solution down, but my timings show little difference when I applied that change. So your speedup must be coming from neatly packing things into the new functionality, I agree that it is cleaner :-) Jul 6 at 19:06